:: Library Catalog

Cover Image

Saved in:

Bibliographic Details
Main Authors:	Cabezas, Alberto, Corenflos, Adrien, Lao, Junpeng, Louf, Rémi, Carnec, Antoine, Chaudhari, Kaustubh, Cohn-Gordon, Reuben, Coullon, Jeremie, Deng, Wei, Duffield, Sam, Durán-Martín, Gerardo, Elantkowski, Marcin, Foreman-Mackey, Dan, Gregori, Michele, Iguaran, Carlos, Kumar, Ravin, Lysy, Martin, Murphy, Kevin, Orduz, Juan Camilo, Patel, Karm, Wang, Xi, Zinkov, Rob
Format:	Preprint
Published:	2024
Subjects:	Mathematical Software Machine Learning Computation
Online Access:	https://arxiv.org/abs/2402.10797
Tags:	Add Tag No Tags, Be the first to tag this record!

Similar Items

JAXbind: Bind any function to JAX
by: Roth, Jakob, et al.
Published: (2024)

DrJAX: Scalable and Differentiable MapReduce Primitives in JAX
by: Rush, Keith, et al.
Published: (2024)

DAC-JAX: A JAX Implementation of the Descript Audio Codec
by: Braun, David
Published: (2024)

MDPax: GPU-acclerated MDP solvers in Python with JAX
by: Farrington, Joseph, et al.
Published: (2025)

laplax -- Laplace Approximations with JAX
by: Weber, Tobias, et al.
Published: (2025)

MPAX: Mathematical Programming in JAX
by: Lu, Haihao, et al.
Published: (2024)

Orbax: Distributed Checkpointing with JAX
by: Gaffney, Colin, et al.
Published: (2026)

Automatic Functional Differentiation in JAX
by: Lin, Min
Published: (2023)

MPX: Mixed Precision Training for JAX
by: Gräfe, Alexander, et al.
Published: (2025)

gyaradax: Local Gyrokinetics JAX Code
by: Galletti, Gianluca, et al.
Published: (2026)

CAX: Cellular Automata Accelerated in JAX
by: Faldor, Maxence, et al.
Published: (2024)

JAXtronomy: A JAX port of lenstronomy
by: Huang, Alan, et al.
Published: (2026)

Optimistix: modular optimisation in JAX and Equinox
by: Rader, Jason, et al.
Published: (2024)

SNNAX -- Spiking Neural Networks in JAX
by: Lohoff, Jamie, et al.
Published: (2024)

minimax: Efficient Baselines for Autocurricula in JAX
by: Jiang, Minqi, et al.
Published: (2023)

Crazyflow: An Accurate, GPU-Accelerated, Differentiable Drone Simulator in JAX
by: Schuck, Martin, et al.
Published: (2026)

rodeo: Probabilistic Methods of Parameter Inference for Ordinary Differential Equations
by: Wu, Mohan, et al.
Published: (2025)

TORAX: A Fast and Differentiable Tokamak Transport Simulator in JAX
by: Citrin, Jonathan, et al.
Published: (2024)

PuzzleJAX: A Benchmark for Reasoning and Learning
by: Earle, Sam, et al.
Published: (2025)

frax: Fast Robot Kinematics and Dynamics in JAX
by: Morton, Daniel, et al.
Published: (2026)

Kozax: Flexible and Scalable Genetic Programming in JAX
by: de Vries, Sigur, et al.
Published: (2025)

Python-JAX-based Fast Stokesian Dynamics
by: Torre, Kim William, et al.
Published: (2025)

NAVIX: Scaling MiniGrid Environments with JAX
by: Pignatelli, Eduardo, et al.
Published: (2024)

Madsen-s-research-group/msmJAX: Initial release
by: Ralf Wanzenböck
Published: (2025)

CLAX: Fast and Flexible Neural Click Models in JAX
by: Hager, Philipp, et al.
Published: (2025)

JaxSGMC: Modular stochastic gradient MCMC in JAX
by: Thaler, Stephan, et al.
Published: (2025)

Chargax: A JAX Accelerated EV Charging Simulator
by: Ponse, Koen, et al.
Published: (2025)

JAXMg: A multi-GPU linear solver in JAX
by: Wiersema, Roeland
Published: (2026)

msmJAX: Fast and Differentiable Electrostatics on the GPU in Python
by: Buchner, Florian, et al.
Published: (2025)

Foragax: An Agent-Based Modelling Framework Based on JAX
by: Chaturvedi, Siddharth, et al.
Published: (2024)

JPC: Flexible Inference for Predictive Coding Networks in JAX
by: Innocenti, Francesco, et al.
Published: (2024)

squishyplanet: Modeling Transits of Non-spherical Exoplanets in JAX
by: Cassese, Ben, et al.
Published: (2024)

Fast and Scalable Inference for Spatial Extreme Value Models
by: Chen, Meixi, et al.
Published: (2021)

Plant-Capture Methods for Estimating Population Size from Uncertain Plant Captures
by: Wang, Yiran, et al.
Published: (2024)

Rapid Scaling of Compositional Uncertainty from Sample to Population Levels
by: Wang, Yiran, et al.
Published: (2025)

TORAX: A Fast and Differentiable Tokamak Transport Simulator in JAX
by: Citrin, Jonathan
Published: (2025)

Self-Scaled Broyden Family of Quasi-Newton Methods in JAX
by: Bioli, Ivan, et al.
Published: (2026)

ABMax: A JAX-based Agent-based Modeling Framework
by: Chaturvedi, Siddharth, et al.
Published: (2025)

lrux: Fast low-rank updates of determinants and Pfaffians in JAX
by: Chen, Ao, et al.
Published: (2026)

unxt: A Python package for unit-aware computing with JAX
by: Starkman, Nathaniel, et al.
Published: (2026)