Reverse Dependencies of jax

The following projects have a declared dependency on jax:

• jaxdifferentialgeometry — Differential geometry using jax
• jaxDiversity — jax implementation for metalearning neuronal diversity
• 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
• 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
• jaxio — Input pipelines for JAX, in JAX
• jaxite — A Homomorphic Encryption implementation (CGGI) written in JAX
• jaxkan — A JAX-based implementation of Kolmogorov-Arnold Networks
• JaxKern — Kernels in Jax.
• jaxkern-nightly — Kernels in Jax.
• jaxknife — no summary
• jaxkuramoto — JAX implementation of Kuramoto model.
• jaxlayerlumos — A JAX-based efficient transfer-matrix method framework for optical simulations
• jaxlem — LEM slope stability analysis using JAX
• jaxley — Differentiable neuron simulations.
• jaxley-mech — A Jaxley-based library of ion channels and syanpses for biophysical neuron models.
• jaxlie — Matrix Lie groups in JAX
• jaxlinop-nightly — A JaxLinOp library.
• jaxls — JAX Language Server APIs.
• jaxmarl — Multi-Agent Reinforcement Learning with JAX
• jaxnet — Neural Nets for JAX
• JaxNRSur — Numerical relativity surrogate model for gravitational waveforms in Jax
• jaxns — Nested Sampling in JAX
• jaxonfunctions — Functions for JAX
• jaxonloader — A dataloader, but for JAX/Numpy
• jaxonmodels — JAX models for deep learning
• jaxonnxruntime — Jaxonnxruntime: JAX based ONNX Runtime.
• jaxoplanet — Probabilistic models for exoplanet inference in JAX
• jaxopt — Hardware accelerated, batchable and differentiable optimizers in JAX.
• jaxparrow — Computes the inversion of the cyclogeostrophic balance based on a variational formulation approach, using JAX
• JaxPM — A simple Particle-Mesh implementation in JAX
• jaxpole — A differentiable implementation of an all-pole filter in JAX
• jaxprvis — A visualizer for JaxPrs
• jaxqualin — A package for extracting quasinormal modes from time domain data
• jaxrenderer — Jax implementation of rasterizer renderer.
• jaxrie — Riemannian JAX
• jaxsim — A differentiable physics engine and multibody dynamics library for control and robot learning.
• jaxsnn — jaxsnn is an event-based approach to machine-learning-inspired training and simulation of SNNs, including support for neuromorphic backends (BrainScaleS-2).
• jaxson — A generative art libary based on Jax
• jaxspec — jaxspec is a bayesian spectral fitting library for X-ray astronomy.
• JaxSSO — A framework for structural shape optimization based on automatic differentiation (AD) and the adjoint method, enabled by JAX
• jaxtorch — A jax based nn library
• jaxtpu — A helper package to install the latest JAX on TPUs, along with all necessary dependencies (e.g. libtpu-nightly)
• jaxtra — JAXTRA
• JaxUtils — Utility functions for JaxGaussianProcesses
• jaxutils-nightly — Utility functions for JaxGaussianProcesses
• JaxVidFlow — Video processing library using JAX.
• jaxwell — Jaxwell is JAX + Maxwell: an iterative solver for solving the finite-difference frequency-domain Maxwell equations on NVIDIA GPUs
• jaxwt — Differentiable and gpu enabled fast wavelet transforms in JAX
• jeig — Various eigendecomposition implementations wrapped for jax.
• jeometric — Graph Neural Networks in JAX
• jf1uids — conservative, differentiable, radial 1d fluid solver with astrophysical modules
• jimGW — Gravitatioanl wave data analysis tool in Jax
• jimm — JAX Image Models
• jimmy-vision — Implementations of popular vision models in Jax (Flax NNX)
• jinns — Physics Informed Neural Network with JAX
• jit-env — A Jax interface for Reinforcement Learning environments.
• jkepler — jax tools for modeling Keplerian orbits
• jmp — JMP is a Mixed Precision library for JAX.
• jnkepler — JAX code for modeling nearly-Keplerian orbits
• jnormcorre — Jax-accelerated implementation of normcorre
• jnotype — Jnotype: probabilistic modeling of high-dimensional binary data.
• jo3mnist — no summary
• jouissance — jouissance
• jpu — JAX + Units
• jpviz — Jaxpr Visualisation Tool
• jraph — Jraph: A library for Graph Neural Networks in Jax
• jrenderer — no summary
• jshbtf0302 — State-of-the-art Natural Language Processing for TensorFlow 2.0 and PyTorch
• jsrm — Kinematic and dynamic models of continuum and articulated soft robots.
• jTWA — Truncated Wigner on GPUs based on JAX.
• 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
• jVMC — jVMC: Versatile and performant variational Monte Carlo
• jvt — Vision transformers with JAX & Flax
• kagglefold — Making protein folding accessible in kaggle platform
• kajihs_utils — Fully typed, plausibly practical, and remarkably random utilities for me—and maybe for you too.
• kaljax — Kalman filter implemented in JAX for GPU and CPU usage
• kalmax — Kalman based neural decoding in Jax
• kasearch — KA-Search: Rapid and exhaustive sequence identity search of known antibodies
• kauldron — Kauldron is a ML research library optimized for quick iteration and modularity.
• kernax — Regularized Stein thinning using JAX
• kernelbiome — A kernel-based nonparametric regression and classification framework for compositional data.
• kernex — Stencil computations in JAX.
• kfac-jax — A Jax package for approximate curvature estimation and optimization using KFAC.
• kheperax — A-maze-ing environment in jax
• kima — A data analysis framework for exoplanet detection
• kinjax — This module calculates FK and Jacobian, accelerated by Jax to utilize JIT and GPU parallelization