author
A LaTeX preamble worth reusing
A LaTeX preamble worth reusing
Tick what you want and compile a PDF in your browser. Nothing is sent anywhere β the Typst compiler runs locally as WebAssembly, from the same template and the same data as the pre-built versions.
author
A LaTeX preamble worth reusing
A LaTeX preamble worth reusing

There's more to JAX
The utilities you keep re-writing, collected in one place.

Paths through phase space
Fit a curve through phase-space points β streams, orbits, trajectories.

Physics-informed neural networks for galactic potentials
Physics-informed neural networks for galactic potentials

Singletons for Equinox
Singletons for Equinox

co-lead developer
Constrain gravitational potentials from stellar stream curvature
A JAX package that fits stream tracks with splines and compares their curvature directly to a potential's accelerations. Co-led with Sirui Wu.

Extras for Diffrax
OOP interfaces and vectorization for Diffrax.

Building blocks for quax classes
Mixins for arithmetic, comparison and array protocols, so a new array type is a few lines.

Multivariate Gaussian KDE
Kernel density estimation in JAX β differentiable, vectorized, bandwidth-tunable.
scipy.special in JAX
Special functions for JAX, following the scipy.special API, filling gaps like zeta at negative integers and the modified Bessel K, Gegenbauer and polylogarithm functions.

Dispatch over plotting backends
Dispatch over plotting backends
![The is_annotated logo: square brackets around a teal block and a purple flag, the shape of Annotated[T, metadata], with a yellow badge carrying a tick.](/highlights/is-annotated.webp)
Check whether a type hint is Annotated
Check whether a type hint is Annotated

dataclasses, for everything
replace, fields and asdict, generalized to any object rather than just dataclasses.

Interoperability between galactic dynamics libraries
Interoperability between galactic dynamics libraries

An immutable mapping, JAX-compatible
An immutable mapping, JAX-compatible
Efficiently get the index-0 element of any iterable
Efficiently get the index-0 element of any iterable

Construct checks for optional dependencies
Construct checks for optional dependencies

lead developer
Galactic dynamics in JAX
Orbit integration, potentials and stream generation β GPU-accelerated and fully differentiable.
maintainer
Multiple dispatch in Python
Multiple dispatch for Python, with type-hint-driven method resolution β the dispatch layer under quax and unxt.

lead developer
Coordinates in JAX
Vectors, frames and transformations β differentiable, and unit-aware via unxt.

Pre-quaxed libraries
Drop-in jax.numpy and friends, already wrapped for dispatch over abstract array types.

lead developer
Units in JAX
Unit-aware quantities that survive jit, grad and vmap.
maintainer
Multiple dispatch in JAX
Custom array-ish types that work with JAX primitives β the substrate the rest of the stack builds on.
lead developer
A standard interface for cosmology libraries
A standard interface for cosmology libraries

Stellar stream tracks
Characterize a stream's track with minimal prior knowledge. Astropy-powered.

lead developer
Constrain macroscopic dark matter with lightning
Constraints on macroscopic dark matter from lightning, on Earth and Jupiter. Released with the PRD paper.

Core developer Β· Coordination Committee Β· Strategic Planning
Astronomy in Python
The community core package for astronomy in Python. The collaboration received the 2025 AAS Lancelot M. Berkeley Prize, a 2023 IOP Publishing Top Cited Paper Award and the 2022 ADASS software prize.