JAX

E95194

JAX is a high-performance numerical computing library for Python that combines NumPy-like APIs with automatic differentiation and just-in-time compilation, widely used for machine learning and scientific computing.

All labels observed (2)

Label Occurrences
JAX canonical 11
JAX (stylized) 1

How this entity was disambiguated

Statements (57)

Predicate Object
instanceOf Python library
numerical computing library
open-source software
compatibleWith Flax
Haiku
NumPy
Optax
Python scientific stack
surface form: SciPy ecosystem

TensorFlow Probability (JAX backend)
developedBy Google
Google Research
documentation https://jax.readthedocs.io
https://jax.readthedocs.io/en/latest/
hasAPIStyle NumPy-like API
hasComponent jax.experimental
jax.lax
NumPy
surface form: jax.numpy

jax.random
implements NumPy
surface form: NumPy API subset

XLA-backed array operations
automatic differentiation primitives
license Apache License 2.0
programmingLanguage Python
repository https://github.com/google/jax
supportsFeature GPU acceleration
TPU acceleration
XLA compilation
automatic differentiation
custom gradients
differentiation of Python functions
forward-mode automatic differentiation
functional transformations
grad-based optimization
higher-order differentiation
jit compilation decorator
just-in-time compilation
just-in-time compiled NumPy operations
just-in-time compiled control flow
parallelization
pmap parallel mapping
random number generation
reverse-mode automatic differentiation
vectorization
vmap vectorized mapping
targetUser engineers
machine learning researchers
scientists
usedFor deep learning
differentiable programming
large-scale linear algebra
machine learning research
neural network training
numerical optimization
probabilistic modeling
scientific computing
simulation-based inference
writtenIn Python

How these facts were elicited

Referenced by (12)

Full triples — surface form annotated when it differs from this entity's canonical label.

NumPy influenced JAX
TPU supportsFramework JAX
Theano influenced JAX
Jax canBeSpelledAs JAX
this entity surface form: JAX (stylized)