Reverse Dependencies of jax

The following projects have a declared dependency on jax:

• jax-autovmap — Automatically broadcast inputs by dynamically applying jax.vmap given un-batched argument dimensions.
• jax-bandflux — A JAX-based package for calculating supernovae Bandfluxes
• jax-bigvgan — JAX Implementation BigVGAN
• jax-causal — Add your description here
• jax-cfd — no summary
• jax-chacha-prng — A pseudo-random number generator for JAX based on the 20 round ChaCha cipher.
• jax_chmc — An implementation of a family of MCMC methods on implicitly defined manifolds
• jax-codex — COders and DEcoders for jaX.
• jax-control-algorithms — Algorithms for state estimation, control, and system identification in JAX
• jax-data — Native data handling for JAX
• jax-dataclasses — Dataclasses + JAX
• jax-dataloader — Dataloader for jax
• jax-dataloaders — A high-performance data loading library for JAX
• jax-ddp — no summary
• jax-dimenet — DimeNet++ in Jax.
• jax-dips — Differentiable 3D interfacial PDE solvers written in JAX using the Neural Bootstrapping Method.
• jax-enums — JAX-compatible Enumerations.
• jax-f16 — Jax F16 dynamics
• jax-fcpe — JAX Implementation FCPE
• jax-fdm — Auto-differentiable and hardware-accelerated force density method
• jax-fid — FID computation in Jax/Flax.
• jax-finufft — JAX bindings for the Flatiron Institute Nonuniform Fast Fourier Transform library
• jax-fixedpoint — Implicit and competitive differentiation in JAX.
• jax-fixedpoint-test-manueldelverme — Implicit and competitive differentiation in JAX.
• jax-flash-attn2 — Flash Attention Implementation with Multiple Backend Support and Sharding This module provides a flexible implementation of Flash Attention with support for different backends (GPU, TPU, CPU) and platforms (Triton, Pallas, JAX).
• jax-flows — Normalizing Flows for JAX
• jax-grid-search — Distributed grid search in JAX
• jax-healpy — Healpix JAX implementation.
• jax-ifnt — Execute runtime assertions, indexing checks, and more if jax code is not traced.
• jax-jumpy — Common backend for Jax or Numpy.
• jax-jynx — A neural network library using jax
• jax_loop_utils — Libraries for ML training loops in JAX.
• jax-lorax — A JAX transform which applies LoRA to arbitrary JAX functions/models
• jax-md — Differentiable, Hardware Accelerated, Molecular Dynamics
• jax-metrics — no summary
• jax-ml-stack — no summary
• jax-nca — Neural Cellular Automata (https://distill.pub/2020/growing-ca/ -- Mordvintsev, et al., "Growing Neural Cellular Automata", Distill, 2020) implemented in JAX
• jax-nerf — Jax implementation of neural radiance fields
• jax-optix — Zero-overhead functional lensing for JAX PyTrees
• jax-quant-finance — High Performance Quantative Finance Library on JAX
• jax-quantity — Quantities in JAX
• jax-radio-tools — Tools used for radio emission of air showers in the JAX framework
• jax-random-projections — sklearn's random projection with JAX to run on a GPU
• jax-relax — JAX-based Recourse Explanation Library
• jax-resnet — Framework-agnostic library for checking array shapes at runtime.
• jax-russell — Financial options valuations and auto-diffed risk measures ("greeks" in Jax.
• jax-scalify — JAX Scalify: end-to-end scaled arithmetic.
• jax-scipy-spatial — Scipy spatial API for JAX
• jax-sgmc — Stochastic Gradient Monte Carlo in Jax
• jax-smi — JAX Synergistic Memory Inspector
• jax-spatial — A jax pytree library for spatial vector algebra
• jax-spcs-kinematics — Implements the Selective Piecewise Constant Strain (SPCS) approach to model the kinematics of continuum soft robots.
• jax-sph — JAX-SPH: Smoothed Particle Hydrodynamics in JAX
• jax-sysid — A Python package for linear and nonlinear system identification and nonlinear regression using Jax.
• jax-toolkit — A collection of jax functions to help with common machine/deep learning related functionality.
• jax-tqdm — Tqdm progress bar for JAX scans and loops
• jax-transformer — Jax Transformer - Jax
• jax-triton — JAX + OpenAI Triton integration
• jax-verify — A library for neural network verification.
• jax-vocos — JAX Implementation Vocos
• jax-xc — no summary
• jax2d — A JAX 2D Rigid-Body Physics Engine
• jax2onnx — export JAX to ONNX - focus on flax nnx
• jax2torch — Jax 2 Torch
• jaxadi — **JaxADi** is a powerful Python library designed to bridge the gap between `casadi.Function` and JAX-compatible functions.
• jaxagents — JAX implementation of Reinforcement Learning agents
• jaxampler — A JAX-based statistical sampling toolkit
• jaxbind — Bind any function written in another language to JAX with support for JVP/VJP/batching/jit compilation
• jaxbook — Iterative algorithms for numerical linear algebra with JAX
• jaxcam — no summary
• jaxcmr — modeling memory search with jax
• jaxcolors — JAX-based color transformation
• jaxdecomp — JAX bindings for the cuDecomp library
• jaxdf — A JAX-based research framework for writing differentiable numerical simulators with arbitrary discretizations
• jaxdifferentialgeometry — Differential geometry using jax
• jaxDiversity — jax implementation for metalearning neuronal diversity
• jaxdpopt — Efficient DP optimization in JAX
• jaxdsp — Fast, differentiable audio processors on the CPU or GPU, with a browser client for real-time control
• jaxert — no summary
• jaxex — A tool for creating science experiments in jax, torch, brax, etc
• JaxFin — JaxFin is a powerful and versatile Python library designed for pricing exotic options using a range of advanced financial techniques.
• jaxfit — Uses JAX for GPU/TPU accelerated nonlinear least-squares curve fitting
• jaxfm — JAX-based implementation of Flow Matching
• jaxformers — 'Attention is all you need' in JAX (Flax)
• jaxfss — JAX/Flax implementation of finite-size scaling
• jaxgcrl — Blazingly fast goal-conditioned environments based on MJX and BRAX.
• jaxgl — An simple JAX graphics library
• jaxGW — Gravitatioanl wave data analysis tool in Jax
• jaxid — A code base for system identification with Jax.
• jaxili — This package provides tools to execute and implement Implicit Likelihood Inference tools in JAX.
• jaximal — A JAX-based PyTree manipulation library
• jaxincell — JAX-in-Cell: Particle-in-Cell code in JAX
• jaxio — Input pipelines for JAX, in JAX
• jaxite — A Homomorphic Encryption implementation (CGGI) written in JAX
• jaxkan — A JAX implementation of Kolmogorov-Arnold Networks
• JaxKern — Kernels in Jax.
• jaxkern-nightly — Kernels in Jax.
• jaxkit — A collection of deep learning architectures, written in jax and flax.
• jaxknife — no summary
• jaxkuramoto — JAX implementation of Kuramoto model.