Google JAX

Google JAX

Tipus	biblioteca Python
Versió inicial	12 desembre 2018
Versió estable	0.4.24 (6 febrer 2024)
Llicència	Llicència Apache, versió 2.0
Equip
Desenvolupador(s)	Peter Hawkins (en) , Matthew Johnson (en) i Jacob VanderPlas (en)
Fonts de codi  Codi font Codi font
Més informació
Lloc web	jax.readthedocs.io…

Google JAX és un marc d'aprenentatge automàtic per transformar funcions numèriques. Es descriu com reunir una versió modificada d'autograd (obtenció automàtica de la funció de gradient mitjançant la diferenciació d'una funció) i XLA de TensorFlow (àlgebra lineal accelerada). Està dissenyat per seguir l'estructura i el flux de treball de NumPy tan de prop com sigui possible i funciona amb diversos marcs existents com TensorFlow i PyTorch. Les funcions principals de JAX són:

1. grau: diferenciació automàtica
2. jit: compilació
3. vmap: vectorització automàtica
4. pmap: programació SPMD

Funció grau

El codi següent mostra la diferenciació automàtica de la funció de graduació .

# imports
from jax import grad
import jax.numpy as jnp

# define the logistic function
def logistic(x): 
 return jnp.exp(x) / (jnp.exp(x) + 1)

# obtain the gradient function of the logistic function
grad_logistic = grad(logistic)

# evaluate the gradient of the logistic function at x = 1 
grad_log_out = grad_logistic(1.0) 
print(grad_log_out)

Funcio jit

El codi següent mostra l'optimització de la funció jit mitjançant la fusió.

# imports
from jax import jit
import jax.numpy as jnp

# define the cube function
def cube(x):
 return x * x * x

# generate data
x = jnp.ones((10000, 10000))

# create the jit version of the cube function
jit_cube = jit(cube)

# apply the cube and jit_cube functions to the same data for speed comparison
cube(x)
jit_cube(x)

Funció vmap

El codi següent mostra la vectorització de la funció vmap.

# imports
from functools import partial
from jax import vmap
import jax.numpy as jnp

# define function
def grads(self, inputs):
 in_grad_partial = partial(self._net_grads, self._net_params)
 grad_vmap = vmap(in_grad_partial)
 rich_grads = grad_vmap(inputs)
 flat_grads = np.asarray(self._flatten_batch(rich_grads))
 assert flat_grads.ndim == 2 and flat_grads.shape[0] == inputs.shape[0]
 return flat_grads

Funció pmap

El codi següent mostra la paral·lelització de la funció pmap per a la multiplicació de matrius.

# import pmap and random from JAX; import JAX NumPy
from jax import pmap, random
import jax.numpy as jnp

# generate 2 random matrices of dimensions 5000 x 6000, one per device
random_keys = random.split(random.PRNGKey(0), 2)
matrices = pmap(lambda key: random.normal(key, (5000, 6000)))(random_keys)

# without data transfer, in parallel, perform a local matrix multiplication on each CPU/GPU
outputs = pmap(lambda x: jnp.dot(x, x.T))(matrices)

# without data transfer, in parallel, obtain the mean for both matrices on each CPU/GPU separately
means = pmap(jnp.mean)(outputs)
print(means)

Biblioteques que utilitzen JAX

Diverses biblioteques de Python utilitzen JAX com a backend, incloent:

• Flax, una biblioteca de xarxes neuronals d'alt nivell desenvolupada inicialment per Google Brain.
• Equinox, una biblioteca que gira al voltant de la idea de representar funcions parametritzades (incloses les xarxes neuronals) com a PyTrees. Va ser creat per Patrick Kidger.
• Diffrax, una biblioteca per a la solució numèrica d'equacions diferencials, com ara equacions diferencials ordinàries i equacions diferencials estocàstiques.
• Optax, una biblioteca per al processament i optimització de gradients desenvolupada per DeepMind.
• Lineax, una biblioteca per resoldre numèricament sistemes lineals i mínims quadrats lineals.
• RLax, una biblioteca per desenvolupar agents d'aprenentatge de reforç desenvolupada per DeepMind.
• jraph, una biblioteca per a xarxes neuronals gràfics, desenvolupada per DeepMind.
• jaxtyping, una biblioteca per afegir anotacions de tipus per a la forma i el tipus de dades ("dtype") de matrius o tensors.