Add GitHub Actions CI

.github/workflows/ci.yml

@@ -0,0 +1,46 @@
+name: CI
+
+on:
+  push:
+    branches: [main]
+  pull_request:
+    branches: [main]
+
+jobs:
+  test:
+    runs-on: ubuntu-latest
+    strategy:
+      matrix:
+        python-version: ["3.11", "3.12", "3.13"]
+    steps:
+      - uses: actions/checkout@v4
+
+      - uses: astral-sh/setup-uv@v5
+        with:
+          enable-cache: true
+
+      - name: Set up Python ${{ matrix.python-version }}
+        run: uv python install ${{ matrix.python-version }}
+
+      - name: Install dependencies
+        run: uv sync --extra dev
+
+      - name: Run tests
+        run: uv run pytest irtk/tests/ -x -q --timeout=60
+        env:
+          JAX_PLATFORMS: cpu
+
+  build:
+    runs-on: ubuntu-latest
+    steps:
+      - uses: actions/checkout@v4
+
+      - uses: astral-sh/setup-uv@v5
+        with:
+          enable-cache: true
+
+      - name: Build package
+        run: uv build
+
+      - name: Check wheel contents
+        run: python3 -m zipfile -l dist/*.whl | grep -c "test_" | xargs -I{} test {} -eq 0

.github/workflows/publish.yml

@@ -0,0 +1,39 @@
+name: Publish to PyPI
+
+on:
+  release:
+    types: [published]
+
+jobs:
+  build:
+    runs-on: ubuntu-latest
+    steps:
+      - uses: actions/checkout@v4
+
+      - uses: astral-sh/setup-uv@v5
+        with:
+          enable-cache: true
+
+      - name: Build package
+        run: uv build
+
+      - name: Check wheel contents
+        run: python3 -m zipfile -l dist/*.whl | grep -c "test_" | xargs -I{} test {} -eq 0
+
+      - uses: actions/upload-artifact@v4
+        with:
+          name: dist
+          path: dist/
+
+  publish:
+    needs: build
+    runs-on: ubuntu-latest
+    permissions:
+      id-token: write
+    steps:
+      - uses: actions/download-artifact@v4
+        with:
+          name: dist
+          path: dist/
+
+      - uses: pypa/gh-action-pypi-publish@release/v1

README.md

@@ -10,12 +10,6 @@ Initially vibe-coded by Opus, so YMMV. PRs welcome.
 pip install irtk-jax
 ```
 
-For Apple Silicon GPU acceleration via [jax-mps](https://github.com/danielpcox/jax-mps):
-
-```bash
-pip install irtk-jax[mps]
-```
-
 ### Development setup
 
 ```bash

irtk/tests/test_model_internal_consistency.py

@@ -43,7 +43,7 @@ def test_logit_lens_values(model_and_tokens):
     model, tokens = model_and_tokens
     result = logit_lens_consistency(model, tokens, position=-1)
     for p in result["per_layer"]:
-        assert -1.0 <= p["logit_cosine"] <= 1.0
+        assert -1.0 - 1e-6 <= p["logit_cosine"] <= 1.0 + 1e-6
     assert result["per_layer"][-1]["agrees_with_final"] is True
 
 
@@ -74,7 +74,7 @@ def test_orthogonality_structure(model_and_tokens):
 def test_orthogonality_values(model_and_tokens):
     model, tokens = model_and_tokens
     result = component_orthogonality(model, tokens, layer=0, position=-1)
-    assert -1.0 <= result["cosine"] <= 1.0
+    assert -1.0 - 1e-6 <= result["cosine"] <= 1.0 + 1e-6
     assert 0 <= result["orthogonality"] <= 1.0
 
 

pyproject.toml

@@ -35,10 +35,6 @@ dev = [
     "pytest>=7.0",
     "torch>=2.0",
 ]
-mps = [
-    "jax-mps",
-]
-
 [tool.hatch.build.targets.wheel]
 packages = ["irtk"]
 exclude = ["irtk/tests"]
@@ -49,12 +45,9 @@ exclude = ["notebooks", "irtk/tests"]
 [tool.pytest.ini_options]
 testpaths = ["irtk/tests"]
 
-[tool.uv.sources]
-jax-mps = { path = "../jax-mps", editable = true }
-
 [dependency-groups]
 dev = [
-    "jax-mps",
     "jupyterlab>=4.5.5",
+    "pytest-timeout>=2.2.0",
     "pytest-xdist>=3.8.0",
 ]