Skip to article frontmatterSkip to article content
Site not loading correctly?

This may be due to an incorrect BASE_URL configuration. See the MyST Documentation for reference.

jaxstro

Evidence-first JAX infrastructure for differentiable science

jaxstro is the evidence-first foundation under the differentiable astrophysics ecosystem and a compact scientific JAX standard library. The active sibling packages gravax, progenax, fluxax, and startrax build on this foundation; stellax remains planned. jaxstro therefore keeps its own dependency surface small.

Astronomy is the proving ground. The reusable product is broader: generic, differentiable, dependency-light scientific infrastructure with explicit units, clear boundary behavior, and validation hooks that make numerical claims auditable. jaxstro owns the primitives that should be shared across scientific packages: physical constants in CGS, unit systems, coordinate transforms, AD-safe interpolation and integration, local basis functions, small dense linear algebra, sampling helpers, parameter-vector bridges, and trust reports. It deliberately refuses to absorb solver libraries; those belong one layer up (see Choose your route, Astro-first, science-general vision, and ADR 0001 — Thin foundation posture).

jaxstro.units is the current canonical ecosystem contract. jaxstro.quantity is implemented and available for evaluation, but ecosystem adoption and any replacement cutover are deferred. Nothing on this site implies that downstream packages should migrate now.

This site is the package’s single source of truth. It is written for two readers at once: a new graduate student meeting differentiable scientific computing for the first time, and future-you trying to remember why a function rounds the way it does. The theory pages teach the methods; the reference pages let you look up the call signature; the decision log records every choice and its trade-offs. Research students and course learners should begin with How to learn with Jaxstro: predict, compute, audit, which explains the site’s recurring predict → compute → audit reasoning cycle. The optional Foundations: the ideas we will not assume route rebuilds the connected ideas the site will not assume are currently active.

Choose your route

Scientific capabilities

Astro-first, not astro-only

The Science patterns enabled by Jaxstro page starts from recurring research questions and routes each one to its generic primitive, transform contract, validation evidence, and downstream ownership boundary.

jaxstro is useful outside astronomy wherever a project needs JAX-native scientific primitives whose gradient behavior is part of the contract. If you are building differentiable models in physics, geoscience, instrumentation, engineering, statistics, or simulation-adjacent inference, the same foundation rules apply: name units and domains, avoid hidden Python-side state, keep shapes static where JAX needs them, validate automatic differentiation against independent checks, and separate generic numerics from domain interpretation.

The long-term vision is not to become a replacement for SciPy, NumPyro, Diffrax, or domain simulators. It is to be a small, rigorous base layer that makes those larger systems easier to trust when an astronomy package, or any differentiable science package, needs shared constants, transformations, numerical kernels, and evidence.

Routed paths

What jaxstro is not

It is not a simulation package, not an inference framework, not a generic SciPy clone, and not an all-purpose solver stack. It holds no domain physics — no stellar tracks, no N-body integrator, no IMF, no filters or photometry semantics. Those live in the packages above it. jaxstro stays at the bottom of the dependency graph and stays light, because everything depends on it.