**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.]])
foo(mu)
# gradient calculation
@tf.function
def grad(mu):
with tf.GradientTape(watch_accessed_variables=False) as t:
t.watch(mu)
property = foo(mu)
loss = t.gradient(property,mu)
return(loss)
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:
t.watch(mu)
property = foo(mu)
print('tracing')
tf.print('executing')
loss = t.gradient(property,mu)
hess = acc.jvp(loss)
return(hess)
tangents = np.zeros(mu.shape)
tangents[0]=1
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)
17
18
1 frames
/usr/local/lib/python3.7/dist-packages/tensorflow/python/framework/func_graph.py 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
```