Determine if in a GradientTape context

General Discussion
, ,
#1

Hi all, is there a way to determine from within an op if the op has been invoked under the context of a GradientTape?

Something like the following:

@tf.function
def some_op():
    gt = get_gradient_tape_somehow()
    if gt:
        # we are in a gradient tape context
        return 1
    return 0

strategy = SomeDistributionStrategy()
with strategy.scope():
    with GradientTape() as tape:
        out = strategy.run(some_op)
#3

Hi @Jack_Hunt

Welcome to the TensorFlow Forum!

You need to put the op inside the TF Gradient Tape to get it recorded. There are few ops which does not have gradient registered. You can check if the op is inside the GradientTape registered by using below code:

image = tf.Variable([[[0.5, 0.0, 0.0]]])
delta = tf.Variable(0.1)

with tf.GradientTape() as tape:
    new_image = tf.image.adjust_contrast(image, delta)
try:
  print(tape.gradient(new_image, [image, delta]))
  assert False   # This should not happen.
except LookupError as e:
  print(f'{type(e).__name__}: {e}')
#Output:
#LookupError: gradient registry has no entry for: AdjustContrastv2