Help protect the Great Barrier Reef with TensorFlow on Kaggle Join Challenge

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 index_depth that matches tf.rank(tensor):

tensor = [[1, 1], [1, 1], [1, 1]]    # tf.rank(tensor) == 2
indices = [[0, 1], [2, 0]]           # num_updates == 2, index_depth == 2
updates = [5, 10]                    # num_updates == 2
print(tf.tensor_scatter_nd_update(tensor, indices, updates))
tf.Tensor(
    [[ 1  5]
     [ 1  1]
     [10  1]], shape=(3, 2), dtype=int32)

Slice updates

When the input tensor has more than one axis scatter can be used to update entire slices.

In this case it's helpful to think of the input tensor as being a two level array-of-arrays. The shape of this two level array is split into the outer_shape and the inner_shape.

indices indexes into the outer level of the input tensor (outer_shape). and replaces the sub-array at that location with the corresponding item from the updates list. The shape of each update is inner_shape.

When updating a list of slices the shape constraints are:

num_updates, index_depth = indices.shape.as_list()
inner_shape = tensor.shape[:index_depth]
outer_shape = tensor.shape[index_depth:]
assert updates.shape == [num_updates, inner_shape]

For example, to update rows of a (6, 3) tensor:

tensor = tf.zeros([6, 3], dtype=tf.int32)

Use an index depth of one.

indices = tf.constant([[2], [4]])     # num_updates == 2, index_depth == 1
num_updates, index_depth = indices.shape.as_list()

The outer_shape is 6, the inner shape is 3:

outer_shape = tensor.shape[:index_depth]
inner_shape = tensor.shape[index_depth:]

2 rows are being indexed so 2 updates must be supplied. Each update must be shaped to match the inner_shape.

# num_updates == 2, inner_shape==3
updates = tf.constant([[1, 2, 3],
                       [4, 5, 6]])

Altogether this gives:

tf.tensor_scatter_nd_update(tensor, indices, updates).numpy()
array([[0, 0, 0],
       [0, 0, 0],
       [1, 2, 3],
       [0, 0, 0],
       [4, 5, 6],