A tensor is a container for data, usually numbers. A matrix is a tensor with two axes; tensors can also have no axes, one axis, or several. In JAX, the numerical array object is called an array. JAX provides NumPy-style array operations through `jax.numpy`, commonly imported as `jnp`. section 2.2 section 3.5.2
jnp$ones(shape = shape(2, 2)) Array([[1., 1.], [1., 1.]], dtype = float32) jnp$zeros(shape = shape(2, 2)) Array([[0., 0.], [0., 0.]], dtype = float32) jnp$array(c(1, 2, 3), dtype = "float32") Array([1., 2., 3.], dtype = float32)
Create a 2 × 2 array of ones, a 2 × 2 array of zeros, and a three-entry array with an explicit `float32` dtype. section 3.5.2
Three attributes describe a tensor: **rank** is its number of axes, **shape** gives the size of each axis, and **dtype** gives the type of its values. A scalar has rank 0 and shape `()`; a five-entry vector has rank 1 and shape `(5)`. A five-entry vector is not a rank-5 tensor: it has five entries along just one axis. For tensor objects, `op_ndim(x)`, `op_shape(x)`, and `op_dtype(x)` inspect these attributes. section 2.2.1 section 2.2.2 section 2.2.5
Handwritten digits as a tensor
A collection of 60,000 grayscale digit images has shape `(60000, 28, 28)` and dtype `uint8`. Its three axes identify the image, its row, and its column. To feed the images to Dense layers, reshape them to `(60000, 784)`, convert them to `float32`, and divide by 255. The data is still numerical, but its shape and value range now suit the model input. section 2.2.5 section 2.5
Shape reflects what each axis means. A batch of feature vectors has shape `(samples, features)`; sequences add a time axis, `(samples, timesteps, features)`. Image batches commonly add height, width, and channels, while video batches add a frames axis. Slicing selects particular elements or regions: for example, selecting images 10 through 99 from the digit data produces a tensor of shape `(90, 28, 28)`. When working from R, check whether slicing uses one-based `@r[]` or zero-based `@py[]` indexing. section 2.2.8 section 2.2.6
Neural-network layers transform tensors using operations such as matrix multiplication, addition, and `relu`. Their weights are tensors too. JAX also introduces gradient functions and just-in-time compilation for numerical computations. Its NumPy-style API has differences in random-number generation and array assignment, so do not assume every NumPy idiom carries over unchanged. section 2.3 section 3.5.1 section 3.5.2
