tf.py_function throws KeyError when computing Jacobian matrix

[trickiwoo]

Click to expand!

Issue Type

Bug

Have you reproduced the bug with TF nightly?

Yes

Source

source

Tensorflow Version

2.11

Custom Code

Yes

OS Platform and Distribution

No response

Mobile device

No response

Python version

No response

Bazel version

No response

GCC/Compiler version

No response

CUDA/cuDNN version

No response

GPU model and memory

No response

Current Behaviour?

tf.py_function throws KeyError when computing Jacobian matrix. It works well when it is computing gradient.

Standalone code to reproduce the issue

import os
import tensorflow as tf
import numpy as np

a = tf.Variable([1.0, 2.0])

def square(a):
    return tf.py_function(lambda a: a ** 2, [a], a.dtype)

with tf.GradientTape(persistent=True) as tape:
    tape.watch(a)
    y = square(a)
# gradient = tape.gradient(y, a) # pass
jacobian = tape.jacobian(y, a)
print(jacobian)

Relevant log output

Node: 'gradient_tape/EagerPyFunc'
KeyError: b'pyfunc_0'


         [[{{node gradient_tape/EagerPyFunc}}]] [Op:__inference_f_173]

[synandi]

Hi @trickiwoo, Apologies for the delay. We were able to replicate the issue in Colab using TF v2.11. Please find the gist here. It seems like we have to dig deep into the issue, we'll update here soon. Thank you!