Add GitHub CI for Python and Octane tests

Python tests run on 3.14 with pytest -n auto. Octane tests run
cargo test --release, skipping airbus tests (need PTY).

Golden test tolerance: strict by default (rtol=1e-4), CI uses
HIGHJAX_TESTS_LOOSE=1 for loose tolerance (rtol=1e-2) to handle
CPU/GPU numerical differences.

Author: cool-RR

.github/workflows/tests.yml

@@ -0,0 +1,66 @@
+name: Tests
+
+on:
+  push:
+    branches: [master]
+  pull_request:
+    branches: [master]
+
+jobs:
+  python-tests:
+    runs-on: ubuntu-latest
+    strategy:
+      matrix:
+        python-version: ['3.14']
+    steps:
+      - uses: actions/checkout@v4
+
+      - name: Set up Python ${{ matrix.python-version }}
+        uses: actions/setup-python@v5
+        with:
+          python-version: ${{ matrix.python-version }}
+          allow-prereleases: true
+
+      - name: Install dependencies
+        run: |
+          pip install -e ".[trainer,tests]"
+
+      - name: Run Python tests
+        env:
+          HIGHJAX_TESTS_LOOSE: '1'
+        run: |
+          pytest -n auto
+
+  octane-tests:
+    runs-on: ubuntu-latest
+    steps:
+      - uses: actions/checkout@v4
+
+      - name: Set up Rust
+        uses: dtolnay/rust-toolchain@stable
+
+      - name: Cache Rust build
+        uses: actions/cache@v4
+        with:
+          path: |
+            octane/target
+            ~/.cargo/registry
+            ~/.cargo/git
+          key: ${{ runner.os }}-cargo-${{ hashFiles('octane/Cargo.lock') }}
+          restore-keys: |
+            ${{ runner.os }}-cargo-
+
+      - name: Set up Python for test data
+        uses: actions/setup-python@v5
+        with:
+          python-version: '3.14'
+          allow-prereleases: true
+
+      - name: Install Python package (for test data generation)
+        run: |
+          pip install -e ".[trainer]"
+
+      - name: Run Octane tests
+        run: |
+          cd octane
+          cargo test --release -- --skip airbus

tests/test_highjax_trainer/test_golden_runs/common.py

@@ -1,8 +1,14 @@
 '''Shared logic for golden run tests.'''
 from __future__ import annotations
 
+import os
+
 import jax.numpy as jnp
 
+_LOOSE = os.environ.get('HIGHJAX_TESTS_LOOSE') == '1'
+_RTOL = 1e-02 if _LOOSE else 1e-04
+_ATOL = 1e-05 if _LOOSE else 1e-07
+
 
 def print_golden_data(epoch_metrics: list[dict]) -> None:
     n = len(epoch_metrics)
@@ -38,7 +44,7 @@ def check_golden_run(epoch_metrics: list[dict], golden_data: dict) -> None:
             expected = expected_values[i]
             actual = float(metrics[key])
 
-            if not jnp.isclose(actual, expected, rtol=1e-04, atol=1e-07):
+            if not jnp.isclose(actual, expected, rtol=_RTOL, atol=_ATOL):
                 bad_fields.append((key, actual, expected))
 
         if bad_fields: