Reverse Dependencies of jax

The following projects have a declared dependency on jax:

• git-t5 — Open source machine learning framework for training T5 models on source code in JAX/Flax.
• glworia — A package for wave-optics lensing calculations
• gmmx — A minimal implementation of Gaussian Mixture Models in Jax
• gnep-learn — gnep-learn - A Python package for learning-based solutions of generalized Nash equilibrium problems.
• goal-jax — (G)eometric (O)ptimiz(A)tion (L)ibraries
• google-jetstream — JetStream is a throughput and memory optimized engine for LLM inference on XLA devices, starting with TPUs (and GPUs in future -- PRs welcome).
• google-vizier — Open Source Vizier: Distributed service framework for blackbox optimization and research.
• google-vizier-dev — Open Source Vizier: Distributed service framework for blackbox optimization and research.
• gotranx — A declarative language describing ordinary differential equations
• gpax — Gaussian processes in NumPyro and JAX
• gpcm — Implementation of the GPCM and variations
• gpfy — Gaussian process with spherical harmonic features in JAX
• GPJax — Gaussian processes in JAX.
• gpjax-nightly — Didactic Gaussian processes in Jax.
• gplugins — gdsfactory plugins
• gpmicrobiome — Temporal probabilistic modeling of bacterial compositions derived from 16S rRNA sequencing
• grad-info-opt — Implementation of Gradient Information Optimization for efficient and scalable training data selection
• GradientTransformation — A JAX-based L-BFGS optimizer
• GradVAR — Gradient update Vector Autoregression modeling library
• grain — Grain: A library for loading and transforming data for ML training.
• grain-nightly — Grain: A library for loading and transforming data for ML training.
• gramform — Grammar for string-to-function formulae
• grapevine-mcmc — Faster Hamiltonian trajectories for problems with guesses, using Blackjax
• graphiq — GraphiQ is a Python library for the simulation, design, and optimization of quantum photonic circuits.
• graphite-datasets — tensorflow/datasets is a library of datasets ready to use with TensorFlow.
• gravlax — Basic training utils for JAX.
• grgrjax — Some generic tools for JAX
• grgrlib — Various insanely helpful functions
• GridPolator — Interpolate a grid of spectroscopic models.
• grins — Gene Regulatory Interaction Network Simulator - GRiNS
• gsmvi — Implementation of Gaussian score matching for variational inference (arXiv:2307.07849)
• gt4py — Python library for generating high-performance implementations of stencil kernels for weather and climate modeling from a domain-specific language (DSL)
• gtech-optimus — Optimus library for real-time marketing personalization using RL.
• gunpowder — A library to facilitate machine learning on large, multi-dimensional images.
• gwkokab — A JAX-based gravitational-wave population inference
• gwpopulation — Unified population inference
• gwpopulation-pipe — A pipeline population inference
• gymnasium — A standard API for reinforcement learning and a diverse set of reference environments (formerly Gym).
• gymnax — JAX-compatible version of Open AI's gym environments
• haiku-geometric — no summary
• haiku-mup — A simple port of μP to Haiku/JAX.
• halospec — Halo Spectroscopy for JWST/MIRI
• hamux — A Deep Learning framework built around ENERGY
• hantoshi — A library to execute PyTorch on TPU
• harmonic — Python package for efficient Bayesian evidence computation
• harmonix — Analytic interferometry of stellar surfaces using spherical harmonics in Jax
• harpy-analysis — single-cell spatial proteomics analysis that makes you happy
• hbmep — Hierarchical Bayesian estimation of motor-evoked potential (MEP) size recruitment curves
• hedeut — HElpful DEcorators and UTilities for working with JAX.
• hedging-gym — HedgingGym is a reinforcement learning environment for training and testing hedging strategies.
• helax — Python package for computing helicity amplitudes
• Helx — Interoperate among reinforcement learning libraries with jax, pytorch, gym and dm_env
• HELX-base — Interoperate among reinforcement learning libraries with jax, pytorch, gym and dm_env
• herculens — Auto-differentiable strong lens modelling
• hgan — A package to infer interpretable dynamics from images of a mechanical system.
• hijax — An experiment framework for Haiku and Jax
• hj-reachability — Hamilton-Jacobi reachability analysis in JAX.
• hnoca — Human Neural Organoid Cell Atlas Toolbox
• Hoi — Higher Order Interactions
• holisticai — no summary
• horqrux — Jax-based quantum state vector and noisy simulator.
• hotaru — High performance Optimizer to extract spike Timing And cell location from calcium imaging data via lineaR impUlse
• hssm — Bayesian inference for hierarchical sequential sampling models.
• hubblax — Simulating Transient in the sky and how to observe them
• hugging-gan-test — Testing pip
• hxrate — HX Rate Fitting
• hybridq — Hybrid Simulator for Quantum Circuits
• hydrax — A zero-copy multiprocess dataloader for JAX.
• hyper-jax — Hyperdimensional computing with Jax
• hyper-nn — Easy hypernetworks in Pytorch and Flax
• hypercoil — Differentiable programming for neuroimaging analysis
• hyperfine — no summary
• hypervecs — no summary
• ic-sharpy — SHARPy is a nonlinear aeroelastic analysis package developed at Imperial College London.
• icland — Recreating Google DeepMind's XLand RL environment in JAX
• icomo — This toolbox aims to simplify the construction of compartmental models and the inference of their parameters
• igenerator — Generating random numbers faster than numpy
• ikpls — Improved Kernel PLS and Fast Cross-Validation.
• image-classification-jax — Run image classification experiments in JAX with ViT, resnet, cifar10, cifar100, imagenette, and imagenet.
• imax — Image augmentation library for Jax
• iminuit — Jupyter-friendly Python frontend for MINUIT2 in C++
• IMNN — Using neural networks to extract sufficient statistics from data by maximising the Fisher information
• impt — Auto-diff Estimator of Lensing Perturbations
• imt-benchmark — High-level Interface to Inertial Motion Tracking
• imt-tree_utils — Utilities for working with Pytrees
• infnum — Infinite and infinitesimal numbers for JAX
• inox — Stainless neural networks in JAX
• interpax — Interpolation and function approximation with JAX
• invrs-gym — A collection of inverse design challenges
• invrs-opt — Algorithms for inverse design
• iqa-jax — IQA library for Jax
• jackofalltrades — Tired of complex machine learning libraries? Introducing jackofalltrades, a streamlined Python package designed to make machine learning accessible for everyone.
• jaims — Library for jax based affine-invariant MCMC sampling
• jammer — Library for jax based affine-invariant MCMC sampling
• jamp — JAMP is a Mixed Precision library for JAX. Forked from DeepMind's JMP.
• janus-sim — A JAX Neural Simulator
• jars — Next-generation objects for omics data
• javiche — A JAX wrapper around ceviche to make interoperability easier. In the future it might make sense to update ceviche itself to use JAX internally.
• jax-ai-stack — no summary
• jax-am — GPU-accelerated simulation toolbox for additive manufacturing based on JAX.