In TensorFlow, standard Tensors are immutable data structures and do not support direct element assignment. To maintain mutable state that allows for index-based updates, elements must be stored within a Variable. However, it is important to note that gradients do not flow backwards through assignments made to a Variable.

TensorFlow Tensor Immutability

Arrays in JAX are strictly immutable, meaning their elements cannot be modified in-place via direct assignment. Instead, JAX provides index update operators, such as .at[].set(), which perform the requested modifications by creating and returning a completely new array populated with the updated values.

JAX Array Immutability

To prevent unnecessary memory allocations, the results of an operation can be assigned to a previously allocated tensor using slice notation. For example, if a tensor `Z` is initialized to have the same shape as `Y` using `zeros_like`, executing `Z[:] = X + Y` overwrites the existing values of `Z` with the newly computed elements. This ensures the update occurs in-place, keeping the memory address of `Z` unchanged.

In-Place Tensor Operations via Slice Assignment

Specific elements or contiguous ranges within a tensor can be modified in-place by specifying their indices and assigning new values. By positioning indexing or slicing operations on the left-hand side of an assignment operator, data can be written directly to those specific locations within the tensor.

Claude

Like elements in standard Python lists, the individual elements within a tensor can be accessed using indexing. Tensor indexing is zero-based, meaning the first element along any given axis is located at index $$0$$. It is important to contrast this programmatic convention with formal linear algebra, where vectors typically utilize one-based indexing (starting at subscript $$1$$).

Tensor Indexing

Dive into Deep Learning

Tensors allow for negative indexing, which is a method used to select elements based on their relative distance from the end of an axis. For instance, using an index of $$-1$$ retrieves the final element along that particular dimension, while $$-2$$ retrieves the second-to-last element.

Negative Tensor Indexing

When a single index or slice is provided for a $$k^{	extrm{th}}$$-order tensor (where $$k > 1$$), the selection operation defaults to being applied along axis $$0$$. For example, applying the index [-1] to a matrix returns its last row, while the slice [1:3] returns its second and third rows.

Learn Before

Related

Learn After