Reverse Dependencies of jax

The following projects have a declared dependency on jax:

• fjformer — Embark on a journey of paralleled/unparalleled computational prowess with FJFormer in JAX
• flaim — Flax Image Models
• flarejax — Flexible Neural Networks in JAX
• flash-attention-jax — Flash Attention - in Jax
• flashbax — Flashbax is an experience replay library oriented around JAX. Tailored to integrate seamlessly with JAX's Just-In-Time (JIT) compilation.
• flax — Flax: A neural network library for JAX designed for flexibility
• flax-addons — flax addons
• flax-gated-linear-rnn — GatedLinearRNN Model
• flax-pilot — A Simplistic trainer for Flax
• flax-trainer — Flax Trainer
• flax-vision-models — A repository of Deep Learning models in Flax
• FlaxDiff — A versatile and easy to understand Diffusion library
• flaxmodels — A collection of pretrained models in Flax.
• flaxsr — Super Resolution tools with Jax/Flax
• flora-opt — no summary
• floral — the best neural network library
• flowfunnel — Leveraging Bayesian hierarchical models to diagnose issues and proactively predict user flow stages in conversion funnels.
• flowjax — Easy to use distributions, bijections and normalizing flows in JAX.
• flowMC — Normalizing flow exhanced sampler in jax
• fmmax — Fourier modal method with Jax
• folding — ColabFold modified for AlphaFold2 and local installation
• folx — Forward Laplacian for JAX
• foragax — ABM package in JAX.
• forecat — Keras based Forescast model builder
• fouriax — A jax port of auraloss
• fttjax — Feature Tokenizer + Transformer - JAX
• funprogram — A functional programming framework for python.
• funsor — A tensor-like library for functions and distributions
• galax — Galactic Dynamics in Jax.
• gatv2-eqx — Add your description here
• gaul — no summary
• gdec — Linear decoders for angled grating stimuli
• gdsfactoryplus — GDSFactory+: adds powerful features such as foundry PDKs, simulations, and verification tools like DRC and LVS.
• gemseo-jax — GEMSEO plugin for JAX (jit compilation, autodiff, XLA)
• generax — Generative Models using Jax
• geometricalgebra — A package for conformal geometric algebra
• geometricconvolutions — Package for building Convolutional Neural Networks on images of tensors.
• geometricjax-experimental — basic tests for full geometricjax library
• gfloat — Generic floating point handling in Python
• ghq — Gauss-Hermite quadrature in JAX
• gigalens — Fast strong gravitational lens modeling
• 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.
• 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
• grain — Grain: A library for loading and transforming data for neural network training.
• grain-nightly — Grain: A library for loading and transforming data for neural network 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.
• 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.
• gwax — Gravitational-wave astronomy in JAX
• 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
• 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
• holistic — no summary
• holisticai — no summary
• horqrux — Jax-based quantum state vector simulator.