Skip to content

Multi-backend (JAX / PyTorch / array-API) support [integration branch]#892

Draft
jobovy wants to merge 18 commits into
mainfrom
feat/backends
Draft

Multi-backend (JAX / PyTorch / array-API) support [integration branch]#892
jobovy wants to merge 18 commits into
mainfrom
feat/backends

Conversation

@jobovy

@jobovy jobovy commented Jun 4, 2026

Copy link
Copy Markdown
Owner

Long-lived integration branch for making all of galpy backend-agnostic and differentiable (numpy/scipy, JAX, PyTorch, array-API). This PR establishes the integration branch: the parallel potential-migration series and the autodiff / action-angle / df work all target feat/backends, are reviewed individually as PRs into it, and the whole branch is merged to main at the end (rebased on main after the variational feature branch #891 merges, since differentiable C orbit integration needs its state-transition matrix).

Draft: integration branch, not for merge until the full stack lands.

Foundation (in this PR)

  • PR0galpy.backend package: a data-first array-namespace resolver (get_namespace: explicit xp= > forced default > array namespace > context/global > numpy), use() context manager, set_default_backend(). numpy/Python inputs resolve to plain numpy (byte-identical path); jax/torch select jax.numpy / array-api-compat torch. Pilot-migrates PlummerPotential + IsochronePotential to private-layer dispatch. 70 new tests green on numpy/jax/torch (parity, grad-vs-FD, the d(Φ)/dR == −Rforce identity, jit/vmap, dtype, decorator pass-through, precedence).
  • P1 — productionized: config [backend] section, a force mode (use(..., force=True)) so a whole run can be pinned to a backend, a --backend {numpy|jax|torch} pytest option + autouse fixture (numpy = no-op), and an AST convention guard banning bare numpy.* in migrated compute methods. Verified --backend jax re-runs test_energy_jacobi_conservation (Plummer/Isochrone) green.

Series targeting this branch

  • P2.x — potential namespace-swap, parallel fan-out by family (spherical / disk / halo-bar / ellipsoidal / dissipative now in progress; then scipy.special-backed after Pspecial; wrappers last).
  • Pspecial — pure-backend special functions (native-preferring router + isolated fallbacks).
  • Tracks D/E/F — autodiff wrappers + in-backend integrators, action-angle, df/streams.

Design: galpy_backend_autodiff_discussion.md; project log in the galpy-jax repo.

🤖 Generated with Claude Code

@jobovy jobovy marked this pull request as ready for review June 4, 2026 17:08
@codecov

codecov Bot commented Jun 4, 2026

Copy link
Copy Markdown

Codecov Report

✅ All modified and coverable lines are covered by tests.
✅ Project coverage is 99.92%. Comparing base (27e01fb) to head (d3cccc0).

Additional details and impacted files
@@           Coverage Diff            @@
##             main     #892    +/-   ##
========================================
  Coverage   99.92%   99.92%            
========================================
  Files         225      243    +18     
  Lines       34224    35161   +937     
  Branches      709      709            
========================================
+ Hits        34197    35136   +939     
+ Misses         27       25     -2     

☔ View full report in Codecov by Harness.
📢 Have feedback on the report? Share it here.

🚀 New features to boost your workflow:
  • ❄️ Test Analytics: Detect flaky tests, report on failures, and find test suite problems.

jobovy added a commit that referenced this pull request Jun 4, 2026
…e test, config docs

- setup.py: remove the 'backend' extra (just array-api-compat — useless on its own);
  keep jax/torch extras (each pulls in array-api-compat).
- _optional_deps.py: correct the array-api-compat comment — there is no vendored
  fallback (removed during coverage cleanup); it is required for the non-numpy
  backends (pulled in by the jax/torch extras) and not needed for numpy-only use.
- Rename tests/test_backend_pilot.py -> tests/test_backend.py; fix the build.yml glob
  to tests/test_backend*.py so it still matches (old _*.py glob would have missed it).
- Document the [backend] config section (default=numpy) in installation.rst + galpyrc.

Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
@jobovy jobovy marked this pull request as draft June 5, 2026 01:43
jobovy and others added 18 commits June 10, 2026 09:13
Add galpy.backend: a data-first array-namespace resolver (get_namespace; precedence
explicit xp= > forced default > namespace of the array args > context/global default
> numpy), a use()/set_default_backend() control with a force mode, and a [backend]
config section. numpy/Python inputs resolve to plain numpy so the numpy path is
byte-identical; jax/torch arrays select jax.numpy / array-api-compat torch.

Migrate PlummerPotential and IsochronePotential to the private-layer dispatch
convention (xp = get_namespace(R, z); numpy.<fn> -> xp.<fn>; numpy.pi -> math.pi) as
the reference for the per-family migrations.

Testing: tests/test_backend.py (resolver + reference potentials), tests/test_backend_
conventions.py (AST guard banning bare numpy.* in migrated compute methods), a
--backend pytest option + autouse fixture to re-run under a forced backend, and a
build.yml job running tests/test_backend*.py under jax+torch. jax/torch extras in
setup.py; config docs for [backend].

Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
…unit parser (#903)

Co-authored-by: Claude Opus 4.8 <noreply@anthropic.com>
Co-authored-by: Claude Opus 4.8 <noreply@anthropic.com>
Co-authored-by: Claude Opus 4.8 <noreply@anthropic.com>
…ties (#913)

Co-authored-by: Claude Opus 4.8 <noreply@anthropic.com>
…#895)

Co-authored-by: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
Co-authored-by: Claude Opus 4.8 <noreply@anthropic.com>
…-swap (#896)

Co-authored-by: Claude Opus 4.8 <noreply@anthropic.com>
Co-authored-by: Claude Opus 4.8 <noreply@anthropic.com>
…906)

Co-authored-by: Claude Opus 4.8 <noreply@anthropic.com>
Co-authored-by: Claude Opus 4.8 <noreply@anthropic.com>
… router fixes) (#922)

Co-authored-by: Claude Opus 4.8 <noreply@anthropic.com>
Co-authored-by: Claude Fable 5 <noreply@anthropic.com>
…re) (#928)

Co-authored-by: Claude Fable 5 <noreply@anthropic.com>
…e evaluation (#929)

Co-authored-by: Claude Opus 4.8 <noreply@anthropic.com>
Co-authored-by: Claude Opus 4.8 <noreply@anthropic.com>
Co-authored-by: Claude Opus 4.8 <noreply@anthropic.com>
…on layer (#932)

Co-authored-by: Claude Opus 4.8 <noreply@anthropic.com>
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant