Google JAX
Google JAX — фреймворк машинного обучения для преобразования числовых функций. Представляет объединение измененной версии autograd (автоматическое получение градиентной функции через дифференцирование функции) и TensorFlow's XLA (Ускоренная линейная алгебра (Accelerated Linear Algebra)). Спроектирован таким образом, чтобы максимально соответствовать структуре и рабочему процессу NumPy для работы с различными существующими фреймворками, такими как TensorFlow и PyTorch. Основными функциями JAX являются:

 grad: автоматическое дифференцирование
 jit: компиляция
 vmap: автоматическая векторизация
 pmap: SPMD программирование

grad 

Код представленный ниже демонстрирует функцию автоматического дифференцирования пакета grad.

# 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) 

Код должен напечатать:

0.19661194

jit 

Код представленный ниже демонстрирует функцию оптимизации через слияние пакета jit.

# 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)

Вычислительное время для jit_cube (строка 17) должно быть заметно короче, чем для cube (строка 16). Увеличение значения в строке 7, будет увеличивать разницу.

vmap 

Код представленный ниже демонстрирует функцию векторизации пакета 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 = jax.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 

Изображение в правой части раздела иллюстрирует идея векторизованного сложения.

pmap 

Код представленный ниже демонстрирует распараллеливание для умножения матриц пакета pmap.

# 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)

Последняя строка должна напечатать значенияː

[1.1566595 1.1805978]

Библиотеки, использующие Jax 
Несколько библиотек Python используют Jax в качестве бэкенда, включая:

 Flax — высокоуровневая библиотека для нейронных сетей изначально разработанная Google Brain.
 Haiku — объектно-ориентированная библиотека для нейронных сетей разработанная DeepMind.
 Equinox — библиотека, основанная на идеи представления параметризованных функций (включая нейронные сети) как PyTrees. Она была создана Патриком Кидгером.
 Optax — библиотека для градиентной обработки и оптимизации разработанная DeepMind.
 RLax — библиотека для разработки агентов для обучения с подкреплением, разработанная DeepMind.

См. также 
 NumPy
 TensorFlow
 PyTorch
 CUDA
 Автоматическое дифференцирование
 JIT-компиляция
 Векторизация
 Автоматическое распараллеливание

Примечания

Ссылки 
 Documentationː 
 Colab (Jupyter/iPython) Quickstart Guideː 
 TensorFlow's XLAː  (Accelerated Linear Algebra)
 
 Original paperː 

Машинное обучение
Google