@Kausta: Tensorflow also contains a documentation for higher order gradients with nested tapes in https://www.tensorflow.org/guide/advanced_autodiff#higher-order_gradients, and it is followed (in the same link) by an input gradient penalty example (The gradient of (the magnitude of the gradient with respect to the inputs) with respect to the mo…

sample.py

# R1 regularization (Hypothetical)
with tf.GradientTape() as t2:
  with tf.GradientTape() as t1:
    # Discriminator outputs
    disc_out = forward(w, x)
  # Regular Loss
  fl = loss(disc_out)
  # Gradients with respect to the inputs
  g = t1.gradient(sum(disc_out), x)
  # Mean squared sum of the gradients
  fl += mean(square(g))
# Gradient of (loss + mse of grads wrt input) with respect to model weights
grad = t2.gradient(fl, w)
w -= grad

# Input gradient penalty example from tensorflow documentation
with tf.GradientTape() as t2:
  # The inner tape only takes the gradient with respect to the input,
  # not the variables.
  with tf.GradientTape(watch_accessed_variables=False) as t1:
    t1.watch(x)
    y = layer(x)
    out = tf.reduce_sum(layer(x)**2)
  # 1. Calculate the input gradient.
  g1 = t1.gradient(out, x)
  # 2. Calculate the magnitude of the input gradient.
  g1_mag = tf.norm(g1)
# 3. Calculate the gradient of the magnitude with respect to the model.
dg1_mag = t2.gradient(g1_mag, layer.trainable_variables)