Smoke: JAX

In [1]:

import os
os.environ.setdefault("XLA_PYTHON_CLIENT_PREALLOCATE", "false")
import jax, jax.numpy as jnp
print(jax.devices())
jax.grad(lambda x: jnp.sum(jnp.sort(x) ** 2))(jnp.array([3., 1., 4., 2.]))

[CudaDevice(id=0)]

Array([6., 2., 8., 4.], dtype=float32)