Reverse Dependencies of chex

The following projects have a declared dependency on chex:

• absl-extra — A wrapper to run and monitor absl app.
• acni — A Jax/Optax implementation of Anticorrelated Noise Injection for Improved Generalization
• aigs — no summary
• alpa — Alpa automatically parallelizes large tensor computation graphs and runs them on a distributed cluster.
• alphafold — An implementation of the inference pipeline of AlphaFold v2.0.This is a completely new model that was entered as AlphaFold2 in CASP14 and published in Nature.
• alphafold-colabfold — An implementation of the inference pipeline of AlphaFold v2.3.1. This is a completely new model that was entered as AlphaFold2 in CASP14 and published in Nature. This package contains patches for colabfold.
• alphafold-kagglefold — An implementation of the inference pipeline of AlphaFold v2.0.This is a completely new model that was entered as AlphaFold2 in CASP14 and published in Nature. This package contains patches for kagglefold
• alphapulldown — Pipeline allows massive screening using alphafold
• apax — Atomistic Learned Potential Package in JAX
• ARLBench — Python Boilerplate that contains all the code you need to create a Python package.
• audiotree — Audio data loading and augmentations in JAX
• axlearn — AXLearn
• bax — A flexible trainer interface for Jax and Haiku.
• bayes-jones — Bayesian inference of Jones matrices.
• bfastor — Fast cryo-EM B-factor refinement
• biorobot — The Bio-inspired Robotics Testbed (BRT).
• boax — Boax is a Bayesian Optimization library for JAX.
• bojaxns — Bayesian Optimisation with JAXNS
• catx — Contextual Bandits with Continuous Actions in JAX
• chemise — Wrapper for training flax models
• chemtrain — Machine-learning deep molecular dynamics potentials.
• chromatix — Differentiable computational optics library using JAX!
• cmonge — Extension of the Monge Gap to learn conditional optimal transport maps
• coax — Plug-n-play reinforcement learning with Gymnasium and JAX
• cr-wavelets — Wavelets with JAX
• craftax — An open-world environment for training RL agents
• dataclass-array — Dataclasses that behave like numpy arrays (with indexing, slicing, vectorization).
• deepqmc — Deep-learning quantum Monte Carlo for electrons in real space
• differt — Differentiable Ray Tracing Toolbox for Radio Propagation Simulations
• diffilqrax — Differentiable iLQR algorithm for dynamical systems
• dinosaur-dycore — no summary
• distrax — Distrax: Probability distributions in JAX.
• dm-clrs — The CLRS Algorithmic Reasoning Benchmark.
• dm-pix — PIX is an image processing library in JAX, for JAX.
• dopamax — Reinforcement learning in pure JAX.
• esquilax — JAX multi-agent simulation and ML toolset
• etils — Collection of common python utils
• evosax — JAX-Based Evolution Strategies
• evox — evox
• exciting-environments — Physical differential equations wrapped into Gymnasium environments
• fjformer — Embark on a journey of paralleled/unparalleled computational prowess with FJFormer in JAX
• flashbax — Flashbax is an experience replay library oriented around JAX. Tailored to integrate seamlessly with JAX's Just-In-Time (JIT) compilation.
• gemseo-jax — GEMSEO plugin for JAX (jit compilation, autodiff, XLA)
• 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.
• gymnax — JAX-compatible version of Open AI's gym environments
• helax — Python package for computing helicity amplitudes
• jaims — Library for jax based affine-invariant MCMC sampling
• jax-dimenet — DimeNet++ 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-moseq — no summary
• jax-scalify — JAX Scalify: end-to-end scaled arithmetic.
• jax-tqdm — Tqdm progress bar for JAX scans and loops
• jaxagents — JAX implementation of Reinforcement Learning agents
• jaxformers — 'Attention is all you need' in JAX (Flax)
• jaxmarl — Multi-Agent Reinforcement Learning with JAX
• jit-env — A Jax interface for Reinforcement Learning environments.
• jumanji — A diverse suite of scalable reinforcement learning environments in JAX
• juxai-s2 — JUX is a jax-accelerated engine for Lux-2022.
• juxai2022 — JUX is a jax-accelerated engine for Lux-2022.
• juxtapose — no summary
• kauldron — Kauldron is a ML research library optimized for quick iteration and modularity.
• mapc-dcf — IEEE 802.11 MAP DCF Simulator
• mapc-sim — IEEE 802.11 MAPC (C-SR) simulator
• mctx — Monte Carlo tree search in JAX.
• minimax-lib — Efficient baselines for autocurricula in JAX
• mlorax — MLoRAx is a minimalist library for low-rank adaptation designd to effortlessly enable parameter-efficient training for Transformer-based models.
• modularbayes — Modular Bayesian Inference.
• MOGPJax — Didactic multi-output Gaussian processes in Jax.
• momaland — A standard API for Multi-Objective Multi-Agent Decision making and a diverse set of reference environments.
• moojoco — A unified framework for implementing and interfacing with MuJoCo and MuJoCo-XLA simulation environments.
• moss-rl — A Python library for Reinforcement Learning.
• mubelnet — Deep Bayesian unsupervised decoder networks. Use poisson or multinomial belief networks to cluster count data.
• neuralgcm — no summary
• neuroflex — An advanced neural network framework with interpretability, generalization, robustness, and fairness features
• nextgenjax — A JAX-based neural network library surpassing Google DeepMind's Haiku and Optax
• optax — A gradient processing and optimization library in JAX.
• orbax-checkpoint — Orbax Checkpoint
• orbax-export — Orbax Export
• ott-jax — Optimal Transport Tools in JAX
• paltax — Strong lensing package using jax
• pax3 — A stateful pytree library for training neural networks.
• phyjax2d — 2D physics simulation by Jax
• praxis — Functionalities such as a layers for building neural networks in Jax.
• precondition-opt — Preconditioning optimizers.
• probabilistic-reconciliation — Probabilistic reconciliation of time series forecasts
• protes — Method PROTES (PRobabilistic Optimizer with TEnsor Sampling) for derivative-free optimization of the multidimensional arrays and discretized multivariate functions based on the tensor train (TT) format
• ptyrodactyl — Ptychography through differentiable programming in JAX
• pyrovelocity — A multivariate RNA Velocity model to estimate future cell states with uncertainty using probabilistic modeling with pyro.
• qdax — A Python Library for Quality-Diversity and NeuroEvolution
• queuinx — Queuinx: A library for performance evaluation in Jax
• ramsey — Probabilistic deep learning using JAX
• reinforced-lib — Reinforcement learning library
• rlax — A library of reinforcement learning building blocks in JAX.
• rmsyutls — Utility functions for Ramsey
• scib-metrics — Accelerated and Python-only scIB metrics
• sequel-core — A Continual Learning Framework for both Jax and PyTorch.
• snaxlib — A simple deep learning library for JAX.
• spaceTree — PyPI package for multi-task label transfer from single-cell refrence data to spatial data
• synthax — SynthAX: A Fast Modular Synthesizer in JAX