Google JAX

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

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

# imports from jax import vmap partial import jax.numpy as jnp  # define function def grads(self, inputs):     in_grad_partial = jax.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 

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