Error in Hessian calculation using forward over backward propagation

System information

  • Have I written custom code (as opposed to using a stock example script provided in TensorFlow): Yes
  • TensorFlow installed from (source or binary): both colab CPU and GPU mode
  • TensorFlow version (use command below):2.7.0
  • Python version: Python 3.7.12
  • GPU model and memory: Tesla P100

Describe the current behavior
I am trying to compute second derivative using the method provide by the tutorial using forward over backward propagation to calculate the hessian-vector-product. For memory reasons, we chose not to use the double backward method. In the code, we have nested loop and tf.dynamic_partition. Both the gradient and hessian works in the eager mode but when I try to decorate the function by @tf.function the error appears, and I found it’s because the combination of using tf.dynamic_partition and for loop gives the error.

Additionally, a different error is given when I try to decorate the hvp function. Without decorator, there’s a TypeError , with a decorator, it gives a SystemError

  • Briefly describe your candidate solution(if contributing): I have tried to find a workaround to use tf.gather or use a loop to replace the tf.dynamic_partition, different problems also prompt.

Standalone code to reproduce the issue
I have made a reduced dummy code here and also in colab

import tensorflow as tf
import numpy as np
@tf.function # without the decorator, the function works fine in eager mode
def foo(mu):

  partitions = tf.constant([1, 0, 0])
  points = tf.dynamic_partition(mu, partitions, 2)[0]
  block = points
  # a dummy example of a loop
  for j in tf.range(1): # without this loop, the function works fine
    block =  points
  return block
# dummy input data
mu = tf.constant([[3.,2.,1.],[3.,2.,1.],[3.,2.,1.]])

# gradient calculation 
def grad(mu):

    with tf.GradientTape(watch_accessed_variables=False) as t:
        property = foo(mu)
    loss = t.gradient(property,mu)

gradient = grad(mu)

# hessian vector product
@tf.function # with/without the decorator, different error prompt
def hvp(mu,tangents):
    with tf.autodiff.ForwardAccumulator(mu, tangents) as acc:
        with tf.GradientTape(watch_accessed_variables=False) as t:
            property = foo(mu)
        loss = t.gradient(property,mu)
    hess = acc.jvp(loss)
tangents = np.zeros(mu.shape)
tangents = tf.convert_to_tensor(tangents,dtype=tf.float32)
hess = hvp(mu,tangents)

Other info / logs Include any logs or source code that would be helpful to
diagnose the problem. If including tracebacks, please include the full
traceback. Large logs and files should be attached.

StagingError                              Traceback (most recent call last)
<ipython-input-6-f99c3f212345> in <module>()
     14 tangents[0]=1
     15 tangents = tf.convert_to_tensor(tangents,dtype=tf.float32)
---> 16 hess = hvp(mu,tangents)

1 frames
/usr/local/lib/python3.7/dist-packages/tensorflow/python/framework/ in autograph_handler(*args, **kwargs)
   1127           except Exception as e:  # pylint:disable=broad-except
   1128             if hasattr(e, "ag_error_metadata"):
-> 1129               raise e.ag_error_metadata.to_exception(e)
   1130             else:
   1131               raise

StagingError: in user code:

    File "<ipython-input-4-f99c3f212345>", line 9, in hvp  *
        loss = t.gradient(property,mu)

    SystemError: PyEval_EvalFrameEx returned a result with an error set

Any help with this would be appreciated!