Forward-mode autodifff fails to process input in tf.Variable type when taking higher-order JVPs and throws Internal error

[eeDigitalSeal]

Click to expand!

Issue Type

Bug

Source

binary

Tensorflow Version

tf 2.9

Custom Code

Yes

OS Platform and Distribution

Linux Ubuntu 20.04

Mobile device

No response

Python version

3.8

Bazel version

No response

GCC/Compiler version

No response

CUDA/cuDNN version

No response

GPU model and memory

No response

Current Behaviour?

If the input type is `tf.Variable`, the forward-mode auto differentiation will throw internal error "ValueError: Internal error: Tried to take gradients (or similar) of a variable without handle data".

Standalone code to reproduce the issue

import tensorflow as tf

primal = tf.Variable(1.1)
with tf.autodiff.ForwardAccumulator(primal, tf.constant(1.)) as outer:
  with tf.autodiff.ForwardAccumulator(primal, tf.constant(1.)) as inner:
    primal_out = primal ** tf.constant(3.5)
inner_jvp = inner.jvp(primal_out)
outer_jvp = outer.jvp(inner_jvp)

Relevant log output

raceback (most recent call last):
  File "/Users/xxx/Documents/Github/bug/bessel_k1.py", line 6, in <module>
    primal_out = primal ** tf.constant(3.5)
  File "/Users/xxx/Documents/Github/venv/lib/python3.8/site-packages/tensorflow/python/ops/variables.py", line 1074, in _run_op
    return tensor_oper(a.value(), *args, **kwargs)
  File "/Users/xxx/Documents/Github/venv/lib/python3.8/site-packages/tensorflow/python/ops/resource_variable_ops.py", line 582, in value
    return self._read_variable_op()
  File "/Users/xxx/Documents/Github/venv/lib/python3.8/site-packages/tensorflow/python/ops/resource_variable_ops.py", line 691, in _read_variable_op
    result = read_and_set_handle()
  File "/Users/xxx/Documents/Github/venv/lib/python3.8/site-packages/tensorflow/python/ops/resource_variable_ops.py", line 681, in read_and_set_handle
    result = gen_resource_variable_ops.read_variable_op(
  File "/Users/xxx/Documents/Github/venv/lib/python3.8/site-packages/tensorflow/python/ops/gen_resource_variable_ops.py", line 479, in read_variable_op
    _result = pywrap_tfe.TFE_Py_FastPathExecute(
  File "/Users/xxx/Documents/Github/venv/lib/python3.8/site-packages/tensorflow/python/eager/forwardprop.py", line 208, in _jvp_dispatch
    return _jvp_exact_shapes(op_name, attr_tuple, inputs, outputs, tangents,
  File "/Users/xxx/Documents/Github/venv/lib/python3.8/site-packages/tensorflow/python/eager/function.py", line 2453, in __call__
    return graph_function._call_flat(
  File "/Users/xxx/Documents/Github/venv/lib/python3.8/site-packages/tensorflow/python/eager/function.py", line 1866, in _call_flat
    forward_function, args_with_tangents = forward_backward.forward()
  File "/Users/xxx/Documents/Github/venv/lib/python3.8/site-packages/tensorflow/python/eager/function.py", line 1397, in forward
    forward_function = self._functions.forward(
  File "/Users/xxx/Documents/Github/venv/lib/python3.8/site-packages/tensorflow/python/eager/function.py", line 1129, in forward
    self._forward_and_backward_functions(inference_args, input_tangents))
  File "/Users/xxx/Documents/Github/venv/lib/python3.8/site-packages/tensorflow/python/eager/function.py", line 1280, in _forward_and_backward_functions
    return self._build_functions_for_outputs(
  File "/Users/xxx/Documents/Github/venv/lib/python3.8/site-packages/tensorflow/python/eager/function.py", line 893, in _build_functions_for_outputs
    forward_wrapper = self._wrap_forward_function_with_jvps(
  File "/Users/xxx/Documents/Github/venv/lib/python3.8/site-packages/tensorflow/python/eager/function.py", line 962, in _wrap_forward_function_with_jvps
    gradient_shape, gradient_dtype = default_gradient.shape_and_dtype(
  File "/Users/xxx/Documents/Github/venv/lib/python3.8/site-packages/tensorflow/python/ops/default_gradient.py", line 40, in shape_and_dtype
    raise ValueError("Internal error: Tried to take gradients (or similar) "
ValueError: Internal error: Tried to take gradients (or similar) of a variable without handle data:
Tensor("Placeholder:0", shape=(), dtype=resource)

[tilakrayal]

@gowthamkpr,
I was able to reproduce the issue on tensorflow v2.8, v2.9 and nightly. Kindly find the gist of it here.