tf.gather_nd implementation mismatch with documentation
@edwardyehuang,
Thank you for the issue. Could you please share a reproducible code that supports your statement so that the issue can be easily debugged? Thank you! 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/62728">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/62728">No</a>
 Hi @edwardyehuang ,

It seems you are missing this part of code below.

https://github.com/tensorflow/tensorflow/blob/4dacf3f368eb7965e9b5c3bbdd5193986081c3b2/tensorflow/python/ops/array_ops.py#L5656

The above piece of code converts scalar Tensor passed to `batch_dims` (Acceptable argument for `tf.gather_nd`) into a python integer.Please refer [gist](https://colab.research.google.com/gist/SuryanarayanaY/75f5be6f00758aed7622aa2d6b91a264/62278.ipynb). As tensor is converted to integer in subsequent steps validation done for integer only.

Thanks! 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/62728">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/62728">No</a>

### Issue type

Bug

### Have you reproduced the bug with TensorFlow Nightly?

Yes

### Source

binary

### TensorFlow version

tf 2.15

### Custom code

No

### OS platform and distribution

Ubuntu 22.04

### Current behavior?

There seems to be a mismatch between the implementation of tf.gather_nd and the documentation.

In documentation, the `batch_dims` can be a tensor.
https://github.com/tensorflow/tensorflow/blob/4dacf3f368eb7965e9b5c3bbdd5193986081c3b2/tensorflow/python/ops/array_ops.py#L5650-L5652

But in the implementation below, only the Python integer is supported.
https://github.com/tensorflow/tensorflow/blob/4dacf3f368eb7965e9b5c3bbdd5193986081c3b2/tensorflow/python/ops/array_ops.py#L5657-L5665

https://github.com/tensorflow/tensorflow/blob/4dacf3f368eb7965e9b5c3bbdd5193986081c3b2/tensorflow/python/ops/array_ops.py#L5684-L5687

Enabling or disabling the autograph does not work.

