D2L - Chapter 2.1 Data Manipulation

import jax
from jax import numpy as jnp
x = jnp.arange(12)
x -> Array([ 0,  1,  2,  3,  4,  5,  6,  7,  8,  9, 10, 11], dtype=int32)
x.size -> 12
x.shape -> (12, )
X = x.reshape(3, 4)
X
jnp.zeros((2, 3, 4))
jnp.ones((2, 3, 4))
# Any call of a random function in JAX requires a key to be
# specified, feeding the same key to a random function will
# always result in the same sample being generated
jax.random.normal(jax.random.PRNGKey(0), (3, 4))
jnp.array([[2, 1, 4, 3], [1, 2, 3, 4], [4, 3, 2, 1]])

X[-1], X[1:3]
# JAX arrays are immutable. jax.numpy.ndarray.at index
# update operators create a new array with the corresponding
# modifications made
X_new_1 = X.at[1, 2].set(17)
X_new_1
X_new_2 = X_new_1.at[:2, :].set(12)
X_new_2

jnp.exp(x)
x = jnp.array([1.0, 2, 4, 8])
y = jnp.array([2, 2, 2, 2])
x + y, x - y, x * y, x / y, x ** y
X = jnp.arange(12, dtype=jnp.float32).reshape((3, 4))
Y = jnp.array([[2.0, 1, 4, 3], [1, 2, 3, 4], [4, 3, 2, 1]])
# axis sets which direction to stack, axis = 0, stacks on the row, meaning will have 6 in axis-0
jnp.concatenate((X, Y), axis=0), jnp.concatenate((X, Y), axis=1)

X == Y
X.sum()

a = jnp.arange(3).reshape((3, 1))
b = jnp.arange(2).reshape((1, 2))
a,
a + b

a = jnp.array([3.5])
a
a.item()
float(a) 🙅
int(a) 🙅