Skip to content

Commit b05eac3

Browse files
alxmrsclaude
alxmrs
and
claude
authored
Add an ml extra for SQL gradients with ddx, and two ML demos (#269)
Adds the `ml` extra (ddxdb), which computes `grad`, `vjp` and `jvp` calls in `Context.sql()` and adds `xarray_sql.ad.sql_all`, and two demos in `benchmarks/ml`, causal attention and nanoGPT, each checked against JAX on DataFusion, DuckDB and DataFusion with einfold. Also adds `benchmarks/ml/CONTRIBUTING.md`, an `ml demos` CI job, and two rules in `AGENTS.md`. Co-Authored-By: Claude Opus 5.5 <[email protected]> Claude-Session: https://claude.ai/code/session_01NyN7idACeDMh94qrnFMu6h
1 parent 94b8a5d commit b05eac3

13 files changed

Lines changed: 2540 additions & 49 deletions

File tree

‎.github/workflows/ci.yml‎

Lines changed: 41 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -65,7 +65,7 @@ jobs:
6565
# Use the --no-install-package to only install the dependencies
6666
# but do not yet build the rust library
6767
- name: Install xarray_sql
68-
run: uv sync --dev --no-install-package xarray-sql
68+
run: uv sync --locked --dev --no-install-package xarray-sql
6969
- name: build rust
7070
run: uv run --no-project maturin develop --uv
7171
- name: Install in-process ADBC drivers
@@ -81,4 +81,43 @@ jobs:
8181
- name: Run unit tests
8282
env:
8383
VIRTUAL_ENV: ${{ github.workspace }}/.venv
84-
run: uv run --no-project pytest -v . -m "not integration"
84+
# The ML demos run in their own job, below.
85+
run: uv run --no-project pytest -v . -m "not integration" --ignore=tests/test_ml_demos.py
86+
87+
ml-demos:
88+
# Each ML demo in benchmarks/ml, at its smallest size, checked against
89+
# JAX (tests/test_ml_demos.py).
90+
name: "ml demos"
91+
runs-on: ubuntu-latest
92+
steps:
93+
- uses: actions/checkout@v4
94+
95+
- uses: dtolnay/rust-toolchain@stable
96+
97+
- name: Setup sccache
98+
uses: mozilla-actions/[email protected]
99+
100+
- name: Configure sccache
101+
run: |
102+
echo "SCCACHE_GHA_ENABLED=true" >> $GITHUB_ENV
103+
echo "RUSTC_WRAPPER=sccache" >> $GITHUB_ENV
104+
105+
- name: Cache Cargo
106+
uses: actions/cache@v4
107+
with:
108+
path: ~/.cargo
109+
key: cargo-cache-${{ steps.rust-toolchain.outputs.cachekey }}-${{ hashFiles('Cargo.lock') }}
110+
111+
- uses: astral-sh/setup-uv@v5
112+
with:
113+
python-version: "3.12"
114+
enable-cache: true
115+
116+
- name: Install xarray_sql
117+
run: uv sync --locked --dev --extra test --no-install-package xarray-sql
118+
- name: build rust
119+
run: uv run --no-project maturin develop --uv
120+
- name: Run the demos
121+
env:
122+
VIRTUAL_ENV: ${{ github.workspace }}/.venv
123+
run: uv run --no-project pytest -v tests/test_ml_demos.py

0 commit comments

Comments
 (0)