Introduction

tf.gather_nd is differentiable with respect to the tensor values it selects, even though it looks like an indexing operation. The important distinction is that TensorFlow can define gradients for params, while the integer indices are treated as fixed selectors and are not differentiable.

What gather_nd Does

gather_nd picks values or slices from a tensor based on index positions.

1import tensorflow as tf
2
3params = tf.constant([[1.0, 2.0],
4                      [3.0, 4.0]])
5indices = tf.constant([[0, 1],
6                       [1, 0]])
7
8result = tf.gather_nd(params, indices)
9print(result)  # [2.0, 3.0]

Conceptually, this is a selection operator. It copies specific entries from params into the output.

Why Gradients Exist for params

A function is differentiable in TensorFlow when the framework knows how output changes map back to input changes. For gather_nd, that mapping is simple for params:

• Selected entries in params influence the output.
• Unselected entries do not.

During backpropagation, TensorFlow sends gradients back only to the positions that were gathered. That backward pass behaves like a scatter-add into the original tensor shape.

You can see it directly:

1import tensorflow as tf
2
3params = tf.Variable([[1.0, 2.0],
4                      [3.0, 4.0]])
5indices = tf.constant([[0, 1],
6                       [1, 0]])
7
8with tf.GradientTape() as tape:
9    gathered = tf.gather_nd(params, indices)
10    loss = tf.reduce_sum(gathered)
11
12grad = tape.gradient(loss, params)
13print(grad)

The gradient is nonzero only at the gathered positions. That is enough to make the operation differentiable with respect to params.

Why indices Are Not Differentiable

The index tensor contains discrete integer locations. Small changes to an integer index do not produce small smooth changes in the output the way ordinary differentiation expects.

That is why TensorFlow does not define useful gradients for indices in the normal autodiff sense. The gather positions are treated as control data, not as learnable continuous variables.

So the correct statement is:

• differentiable with respect to gathered values: yes
• differentiable with respect to integer indices: no

This is the same pattern you see in many deep learning operations that use masks, routing, or table lookups.

Repeated Indices Still Work

If the same position is gathered more than once, the backward pass accumulates gradients at that position.

1import tensorflow as tf
2
3params = tf.Variable([10.0, 20.0, 30.0])
4indices = tf.constant([[1], [1]])
5
6with tf.GradientTape() as tape:
7    values = tf.gather_nd(params, indices)
8    loss = tf.reduce_sum(values)
9
10grad = tape.gradient(loss, params)
11print(grad)

The gradient at index 1 is added twice because that source value contributed to the loss twice.

How to Think About It

A good mental model is that gather_nd is like a sparse linear selection operator applied to params. Linear selection is differentiable with respect to the selected numeric values, even if the selection pattern itself is fixed and discrete.

This is why gather operations fit cleanly inside models that select embeddings, route tokens, or pick specific elements from intermediate tensors.

Common Pitfalls

The most common mistake is assuming that if an operation uses indices, nothing about it can be differentiable. The gathered values are still ordinary numeric inputs, and gradients can flow to them.

Another issue is expecting gradients with respect to indices. Integer locations are discrete selectors, so TensorFlow does not treat them as smooth differentiable variables.

People also get confused when repeated indices produce accumulated gradients. That is expected because the same source entry influenced the loss more than once.

Finally, if gradients appear as zeros, confirm that the corresponding tensor entries actually contributed to the loss. Unselected elements naturally receive no gradient.

Summary

• 'tf.gather_nd is differentiable with respect to params, not with respect to integer indices.'
• Backpropagation sends gradients only to the gathered positions.
• The backward pass behaves like a scatter-add into the original tensor shape.
• Repeated indices accumulate gradient contributions.
• Think of gather_nd as differentiable selection over values with a fixed discrete routing pattern.