ML Community Day is November 9! Join us for updates from TensorFlow, JAX, and more Learn more

tf.tensor_scatter_nd_update

"Scatter updates into an existing tensor according to indices.

This operation creates a new tensor by applying sparse updates to the input tensor. This is similar to an index assignment.

# Not implemented: tensors cannot be updated inplace.
tensor[indices] = updates

If an out of bound index is found on CPU, an error is returned.

  • If an out of bound index is found, the index is ignored.
  • The order in which updates are applied is nondeterministic, so the output will be nondeterministic if indices contains duplicates.

This operation is very similar to tf.scatter_nd, except that the updates are scattered onto an existing tensor (as opposed to a zero-tensor). If the memory for the existing tensor cannot be re-used, a copy is made and updated.

In general:

  • indices is an integer tensor - the indices to update in tensor.
  • indices has at least two axes, the last axis is the depth of the index vectors.
  • For each index vector in indices there is a corresponding entry in updates.
  • If the length of the index vectors matches the rank of the tensor, then the index vectors each point to scalars in tensor and each update is a scalar.
  • If the length of the index vectors is less than the rank of tensor, then the index vectors each point to slices of tensor and shape of the updates must match that slice.

Overall this leads to the following shape constraints:

assert tf.rank(indices) >= 2
index_depth = indices.shape[-1]
batch_shape = indices.shape[:-1]
assert index_depth <= tf.rank(tensor)
outer_shape = tensor.shape[:index_depth]
inner_shape = tensor.shape[index_depth:]
assert updates.shape == batch_shape + inner_shape

Typical usage is often much simpler than this general form, and it can be better understood starting with simple examples:

Scalar updates

The simplest usage inserts scalar elements into a tensor by index. In this case, the index_depth must equal the rank of the input tensor, slice each column of indices is an index into an axis of the input tensor.

In this simplest case the shape constraints are:

num_updates, index_depth = indices.shape.as_list()
assert updates.shape == [num_updates]
assert index_depth == tf.rank(tensor)`

For example, to insert 4 scattered elements in a rank-1 tensor with 8 elements.

This scatter operation would look like this:

tensor = [0, 0, 0, 0, 0, 0, 0, 0]    # tf.rank(tensor) == 1
indices = [[1], [3], [4], [7]]       # num_updates == 4, index_depth == 1
updates = [9, 10, 11, 12]            # num_updates == 4
print(tf.tensor_scatter_nd_update(tensor, indices, updates))
tf.Tensor([ 0 9  0 10  11  0  0 12], shape=(8,), dtype=int32)

The length (first axis) of updates must equal the length of the indices: num_updates. This is the number of updates being inserted. Each scalar update is inserted into tensor at the indexed location.

For a higher rank input tensor scalar updates can be inserted by using an