Infinite loop in batch_jacobian for ode
@dv-ai,
I tried to execute the code on tensorflow v2.12 & tf-nightly, and it was executed without an infinite loop as mentioned. Kindly find the gist of it [here](https://colab.research.google.com/gist/tilakrayal/ddea7e0534f153e96e81eebbeb18cc58/untitled1173.ipynb) and also please have a look at an issue with the similar error which is under the developer priority list.  Thank you!
https://github.com/tensorflow/tensorflow/issues/57343 This issue is stale because it has been open for 7 days with no activity. It will be closed if no further activity occurs. Thank you. This issue was closed because it has been inactive for 7 days since being marked as stale. Please reopen if you'd like to work on this further. Are you satisfied with the resolution of your issue?
<a href="https://docs.google.com/forms/d/e/1FAIpQLSfaP12TRhd9xSxjXZjcZFNXPGk4kc1-qMdv3gc6bEP90vY1ew/viewform?entry.85265664=Yes&entry.2137816233=https://github.com/tensorflow/tensorflow/issues/60675">Yes</a>
<a href="https://docs.google.com/forms/d/e/1FAIpQLSfaP12TRhd9xSxjXZjcZFNXPGk4kc1-qMdv3gc6bEP90vY1ew/viewform?entry.85265664=No&entry.2137816233=https://github.com/tensorflow/tensorflow/issues/60675">No</a>

<details><summary>Click to expand!</summary> 
 
 ### Issue Type

Bug

### Have you reproduced the bug with TF nightly?

Yes

### Source

source

### Tensorflow Version

tf 2.7.0

### 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?

When executing the computation of batch_jacobian on an ode with experimental_use_pfor=True, the code is stuck in an infinity loop. 
When experimental_use_pfor=False, the computation working like expected. 

### Standalone code to reproduce the issue

```shell
import tensorflow as tf
import tensorflow_probability as tfp

@tf.function
def ode_fun(t,x, c):
    return x * c


def compute_grads(solver,experimental_use_pfor):
    solution_times = tf.constant([1.0])
    c = tf.Variable(tf.constant([[1.0],[2.0]]))
    with tf.GradientTape(watch_accessed_variables=False, persistent=not experimental_use_pfor) as tape:
        tape.watch(c)
        result = solver.solve(ode_fun, 0.0, tf.constant([[1.0],[2.0]]), solution_times=solution_times, constants={'c':c}).states[0]

    grads = tape.batch_jacobian(result,c,experimental_use_pfor=experimental_use_pfor,parallel_iterations=2)
    return  grads

def execute():
    print(tf.__version__)
    experimental_use_pfor = True
    solver = tfp.math.ode.DormandPrince()
    print(compute_grads(solver,experimental_use_pfor))

execute()
```


### Relevant log output

```shell
The computation never ends.
```
</details>