tf.py_function throws KeyError when computing Jacobian matrix
Hi @trickiwoo, Apologies for the delay. We were able to replicate the issue in Colab using TF v2.11. Please find the gist [here](https://colab.sandbox.google.com/gist/synandi/40e7a6bb108d9a9b97bcbe76ab0c67da/60084.ipynb). It seems like we have to dig deep into the issue, we'll update here soon. Thank you!
<details><summary>Click to expand!</summary> 
 
 ### 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?

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


### Standalone code to reproduce the issue

```shell
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

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


         [[{{node gradient_tape/EagerPyFunc}}]] [Op:__inference_f_173]
```
</details>