PyRank

Jax Python Packages

Python packages with the GitHub topic jax. Sorted by monthly downloads, with stars and forks.
arogozhnikov
einops

Flexible and powerful tensor operations for readable and reliable code (for pytorch, jax, TF and others)

21.5M 10K 404
explosion
thinc

🔮 A refreshing functional take on deep learning, compatible with your favorite libraries

20.9M 3K 293
wandb
wandb

The AI developer platform. Use Weights & Biases to train and fine-tune models, and manage models from experimentation to production.

15.5M 11K 904
keras-team
keras

Deep Learning for humans

11.5M 64K 20K
jax-ml
jax

Composable transformations of Python+NumPy programs: differentiate, vectorize, JIT to GPU/TPU, and more

9.4M 36K 4K
jax-ml
jaxlib

Composable transformations of Python+NumPy programs: differentiate, vectorize, JIT to GPU/TPU, and more

9.2M 36K 4K
google
etils

Collection of eclectic utils for python.

7.1M 260 39
patrick-kidger
jaxtyping

Type annotations and runtime checking for shape and dtype of JAX/NumPy/PyTorch/etc. arrays. https://docs.kidger.site/jaxtyping/

4.7M 2K 99
HIPS
autograd

Efficiently computes derivatives of NumPy code.

3.3M 8K 947
google
orbax-checkpoint

Orbax provides common checkpointing and persistence utilities for JAX users

2.9M 537 104
google
flax

Flax is a neural network library for JAX that is designed for flexibility.

2.5M 7K 846
patrick-kidger
equinox

Elegant easy-to-use neural networks + scientific computing in JAX. https://docs.kidger.site/equinox/

1.4M 3K 216
google
grain

Library for reading and processing ML training data.

1M 790 91
pyro-ppl
numpyro

Probabilistic programming with NumPy powered by JAX for autograd and JIT compilation to GPU/TPU/CPU.

961K 3K 325
tensorflow
tensorflow-datasets

TFDS is a collection of datasets ready to use with TensorFlow, Jax, ...

794K 5K 2K
jax-ml
jax-cuda12-pjrt

Composable transformations of Python+NumPy programs: differentiate, vectorize, JIT to GPU/TPU, and more

668K 36K 4K
jax-ml
jax-cuda12-plugin

Composable transformations of Python+NumPy programs: differentiate, vectorize, JIT to GPU/TPU, and more

666K 36K 4K
google
jaxopt

Hardware accelerated, batchable and differentiable optimizers in JAX.

512K 1K 75
jax-ml
jax-cuda13-pjrt

Composable transformations of Python+NumPy programs: differentiate, vectorize, JIT to GPU/TPU, and more

458K 36K 4K
jax-ml
jax-cuda13-plugin

Composable transformations of Python+NumPy programs: differentiate, vectorize, JIT to GPU/TPU, and more

442K 36K 4K
google
clu

CLU lets you write beautiful training loops in JAX.

436K 370 34
keras-team
keras-hub

Deep Learning for humans

401K 64K 20K
jcmgray
autoray

Abstract your array operations.

378K 174 13
google
lineax

Linear solvers in JAX and Equinox. https://docs.kidger.site/lineax

376K 585 46