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.gatheror use a loop to replace thetf.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