Repository: tensorflow
Issue: Graph/memory corruption involving custom_gradient of functions involving pow
Labels: stat:awaiting tensorflower,type:bug,comp:autograph,TF 2.12

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

Bug

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

Yes

### Source

binary

### Tensorflow Version

2.9.1, 2.11, presumably more

### Custom Code

Yes

### OS Platform and Distribution

Linux Ubuntu 18.04, 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?

I have implemented `custom_gradient`s for `tf.function`s that smooth various Nth roots to avoid the singularity at `x=0`. I expect the value computed by the custom gradient to be used whenever that function is differentiated, and not substituted for any other value.

**HOWEVER** when the following conditions hold:
- input is a complex dtype
- the `GradientTape` is persistent
- the `tf.function` being differentiated invokes `tf.pow` directly (or indirectly through `**` operator)
- the argument is `0`

Then the custom gradient is invoked and the proper value is computed, but then discarded somewhere between returning from my custom gradient and returning from the gradient tape's internal computation.

Note: this problem does not appear when wrapping `sqrt()` as that has its own gradient implementation that presumably avoids `pow()`.

This seems likely related to numerical issues highlighted in tensorflow/tfjs#346, PLUS, some kind of graph/memory corruption, although why/how that would be conditional on the argument is beyond me.

### Standalone code to reproduce the issue

```shell
import tensorflow as tf


@tf.function
@tf.custom_gradient
def cbrt(p):
  third = 1. / 3.
  real_part_nonnegative = tf.math.real(p) >= 0
  signs = tf.cast(tf.where(real_part_nonnegative, 1.0, -1.0), p.dtype)
  p = signs * p
  unsigned_root = p ** third

  @tf.function
  def grad(upstream):
    unsigned_root_plus_eps = unsigned_root + 1e-4
    denom = 3 * unsigned_root_plus_eps * unsigned_root_plus_eps + 1e-4
    droot_dp = (1. / denom)
    grad_root = droot_dp * upstream
    tf.print("custom gradient:", tf.math.real(grad_root), tf.math.imag(grad_root))
    return grad_root

  root = signs * unsigned_root
  return root, grad
  

if __name__ == "__main__":
  tf.config.run_functions_eagerly(True)
  x = tf.Variable([-1e-2, -1e-10, 0, 1e-10, 1e-2], dtype=tf.complex64)
  with tf.GradientTape(persistent=True) as tape:
    cbrt_x = cbrt(x)
  gradients = tape.gradient(cbrt_x, x)
  tf.print("all good:", tf.math.real(gradients), tf.math.imag(gradients))

  tf.config.run_functions_eagerly(False)
  with tf.GradientTape() as tape:
    cbrt_x = cbrt(x)
  gradients = tape.gradient(cbrt_x, x)
  tf.print("all good:", tf.math.real(gradients), tf.math.imag(gradients))

  @tf.function
  def now_again_in_graph_mode(y):
    with tf.GradientTape(persistent=True) as tape:
      cbrt_y = cbrt(y)
    return tape.gradient(cbrt_y, y)
  gradients = now_again_in_graph_mode(x)
  tf.print("chaos reigns:", tf.math.real(gradients), tf.math.imag(gradients))
```


### Relevant log output

```shell
custom gradient: [7.16964388 9905.4209 9997.00098 9905.4209 7.16964388] [0 0 0 0 0]
all good: [7.16964388 9905.4209 9997.00098 9905.4209 7.16964388] [0 0 0 0 0]
custom gradient: [7.16964388 9905.4209 9997.00098 9905.4209 7.16964388] [0 0 0 0 0]
all good: [7.16964388 9905.4209 9997.00098 9905.4209 7.16964388] [0 0 0 0 0]
custom gradient: [7.16964388 9905.4209 9997.00098 9905.4209 7.16964388] [0 0 0 0 0]
chaos reigns: [7.16964388 9905.4209 -nan 9905.4209 7.16964388] [0 0 -nan 0 0]
```
</details>