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.

AI illustration

How this image was made

AI-generated illustration of JAX

This AI-generated illustration was produced by black-forest-labs/FLUX.2-dev (1024x1024) from a prompt written by openai/gpt-oss-120b from the entity's label + description.

Prompt

Generate an image of JAX (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 (3)

Label Occurrences
JAX canonical 32
JAX (stylized) 1
JAX project 1

How this entity was disambiguated

Statements (57)

Predicate Object
instanceOf Python library ⓘ
numerical computing library ⓘ
open-source software ⓘ
compatibleWith Flax ⓘ
Haiku ⓘ
NumPy ⓘ
Optax ⓘ
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 ⓘ
jax.numpy ⓘ
linked to: NumPy

jax.random ⓘ
implements NumPy API subset ⓘ
linked to: NumPy

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 (34)

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 (stylized) ⓘ
linked to: JAX
XLA → integratedWith → JAX ⓘ
Google Cloud TPU v2 → supportsFramework → JAX ⓘ
subject linked to: Google Cloud TPU V2
Google Cloud TPU v3 → supportsFramework → JAX ⓘ
subject linked to: Google Cloud TPU V3
Google Cloud TPU v4 → supports → JAX ⓘ
subject linked to: Google Cloud TPU V4
Flax → basedOn → JAX ⓘ
jax.random → distributedBy → JAX project ⓘ
linked to: JAX
Adam → implementedIn → JAX ⓘ
Adam → implementedIn → JAX ⓘ
Gated Recurrent Unit → implementedIn → JAX ⓘ
subject linked to: GU
NumPyro → basedOn → JAX ⓘ
Trax → builtOn → JAX ⓘ
subject linked to: Trax library