- 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
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
- Briefly describe your candidate solution(if contributing): I have tried to find a workaround to use
tf.gatheror 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) 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=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=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