tensorflow: Subsequent gather_nd Operation leads to error: Tensor.graph is meaningless when eager execution is enabled.

System information

• OS Platform and Distribution (e.g., Linux Ubuntu 16.04): Ubuntu 16.04
• TensorFlow installed from (source or binary): binary
• TensorFlow version (use command below): 2.0.0-alpha0
• Python version: python 3.5
• CUDA/cuDNN version: 10.
• GPU model and memory: TITAN

Describe the current behavior Gradient computation fails with the message: “Tensor.graph is meaningless when eager execution is enabled.”

The gradient becomes an IndexedSlices through the second gather_nd operation. On the IndexedSlices the check in https://github.com/tensorflow/tensorflow/blob/8cb1aa18882be9142a5e4d3977de0a076ce2c791/tensorflow/python/framework/ops.py#L5956-L5957 evaluates to True and in line 5965 the .graph attribute is called which in EagerExecution causes the error mentioned above. https://github.com/tensorflow/tensorflow/blob/8cb1aa18882be9142a5e4d3977de0a076ce2c791/tensorflow/python/framework/ops.py#L5965

A workaround/fix for the problem can be achieved by adding another check to the if clause making sure that op_input is not an IndexedSlices: and not isinstance(op_input, IndexedSlices)

    if (isinstance(op_input, (Operation, _TensorLike)) and
       ((not isinstance(op_input, Tensor)) or type(op_input) == Tensor) and
       (not isinstance(op_input, IndexedSlices))):  # pylint: disable=unidiomatic-typecheck

Describe the expected behavior

Gradients to be computed correctly

Code to reproduce the issue

import tensorflow as tf


class DebugModel(tf.keras.Model):
  def __init__(self):
    super(DebugModel, self).__init__()
    self.dense = tf.keras.layers.Dense(20, activation='relu', input_shape=(5,))

  def __call__(self, input, *args, **kwargs):
    x = self.dense(input)
    ind = [[0], [1], [2]]
    x = tf.gather_nd(x, ind)
    ind = [[0], [1]]
    x = tf.gather_nd(x, ind)
    return x

my_model = DebugModel()

my_optim = tf.keras.optimizers.Adam(1e-3)
with tf.GradientTape() as tape:
  out = my_model(tf.random.uniform((16,5)))
  my_loss = -tf.reduce_mean(out)
  grads = tape.gradient(my_loss, my_model.trainable_weights)
  my_optim.apply_gradients(zip(grads, my_model.trainable_weights))