Tensorflow numpy_function causes errors when using tf.shape
Note: Tensorflow 2.8.4 didn't have this issue. Upgrading to 2.13.1 caused the error.

An equivalent code that doesn't produce any error in 2.13.1:

```
import tensorflow as tf  # 2.13.1
import tensorflow_datasets as tfds  # 4.9.3
from functools import partial


AUTO = tf.data.experimental.AUTOTUNE

dataset_name = "mnist"
dataset = tfds.load(
    dataset_name,
    split="train[:1%]"
)

dataset_len = dataset.reduce(0, lambda x, _: x + 1).numpy()
print("dataset_len", dataset_len)

image_size = [28, 28, 1]
new_image_size = [14, 2, 28, 1]


def prepare_input(
    image, labels,
    image_size
):

    image = tf.image.resize(image, image_size)

    split_shape = tf.shape(image)[1] // 2
    image_reshape = tf.reshape(
        image, [split_shape, 2, image_size[1], 1]
    )

    return {
        "image": image_reshape,
        "label": labels
    }


partial_prepare_input = partial(
    prepare_input,
    image_size=tf.constant(image_size[:2], dtype=tf.int32)
)


dataset = dataset.map(
    lambda sample: partial_prepare_input(
        sample["image"],
        sample["label"]
    ),
    num_parallel_calls=AUTO
)


for x in dataset:
    print(x['image'].shape, x['label'].shape)

``` Hi @basnetr ,

I have tested the attached code snippet it executes fine with Tf2.1.3.1 and Tf2.15 as well. Please refer attached [gist](https://colab.sandbox.google.com/gist/SuryanarayanaY/54ff05c4391f698c7cdb17ffa18dd915/62778.ipynb).

Please check and confirm. Thanks! @SuryanarayanaY Thanks for testing it in Colab, it does seem to run fine there. However, I still get the error on Linux Ubuntu 20.04.6 LTS. Furthermore, this error seems to appear when running on GPU, code also runs without errors on my environment on CPU with:
```
import os
os.environ["CUDA_VISIBLE_DEVICES"] = "-1"
``` Hi @basnetr ,

I have replicated the behaviour with GPU runtime and attached gist here.The problem arises during iteration of the dataset object with GPU runtime. May need to escalate to SME. Thanks! It seems the behaviour is inconsistent.Repeated runs of same code results sometimes success and sometimes error out.

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/62778">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/62778">No</a>
 @SuryanarayanaY Yes, but it'd be nice to have always success. Currently it's mostly error. Reopening as it is stalled by mistake. I have a similar error with the following code snippet (with tf v2.15.1 and v2.16.1 on ubuntu 22.04):
```
import tensorflow as tf


@tf.numpy_function(Tout = (tf.float32))
def read(number):
    img = tf.constant(number, tf.float32, shape = (512, 424))
    img = tf.image.resize(img[..., tf.newaxis], tf.shape(img) // 2, preserve_aspect_ratio = True)
    return img


dataset = tf.data.Dataset.from_tensor_slices([
    0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 0, 1, 2, 3, 4, 5, 6, 7, 8, 9])
dataset = dataset.map(read, tf.data.AUTOTUNE)
dataset = dataset.batch(4, False, tf.data.AUTOTUNE).prefetch(tf.data.AUTOTUNE)

for img in dataset:
    print(img.shape)
```

Executing this results in different outputs, e.g.:
[test_output_1.txt](https://github.com/tensorflow/tensorflow/files/15392117/test_output_1.txt)
[test_output_2.txt](https://github.com/tensorflow/tensorflow/files/15392137/test_output_2.txt)

Additionally the program never terminates and is apparently stuck waiting on something (0% cpu usage)...
Are these problems maybe symptoms of some race condition somewhere?
### Issue type

Bug

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

Yes

### Source

source

### TensorFlow version

2.13.1

### Custom code

Yes

### OS platform and distribution

Linux Ubuntu 20.04.6 LTS

### Mobile device

_No response_

### Python version

3.9.18

### Bazel version

_No response_

### GCC/compiler version

_No response_

### CUDA/cuDNN version

_No response_

### GPU model and memory

_No response_

### Current behavior?

Raises an error. 

Expect: No errors.

### Standalone code to reproduce the issue

```shell
import tensorflow as tf  # v2.13.0-17-gf841394b1b7 2.13.1
import tensorflow_datasets as tfds  # 4.9.3
from functools import partial

AUTO = tf.data.experimental.AUTOTUNE

dataset_name = "mnist"
dataset = tfds.load(
    dataset_name,
    split="train[:1%]"
)

dataset_len = dataset.reduce(0, lambda x, _: x + 1).numpy()
print("dataset_len", dataset_len)

image_size = [28, 28, 1]
new_image_size = [14, 2, 28, 1]


def prepare_input(
    image, labels,
    image_size
):

    image = tf.image.resize(image, image_size)

    split_shape = tf.shape(image)[1] // 2
    image_reshape = tf.reshape(
        image, [split_shape, 2, image_size[1], 1]
    )

    return image_reshape, labels


partial_prepare_input = partial(
    prepare_input,
    image_size=tf.constant(image_size[:2], dtype=tf.int32)
)


def prepare_input_wrapper(sample):
    image = sample["image"]
    labels = sample["label"]

    labels_shape = labels.shape

    image, label = tf.numpy_function(
        func=partial_prepare_input,
        inp=[image, labels],
        Tout=[tf.float32, tf.int64],
        name="numpy_function_1"
    )

    image.set_shape(new_image_size)
    label.set_shape(labels_shape)

    return {
        "image": image,
        "label": label
    }


dataset = dataset.map(
    lambda sample: prepare_input_wrapper(sample),
    num_parallel_calls=AUTO
)


for x in dataset:
    print(x['image'].shape, x['label'].shape)
```


### Relevant log output

```shell
2024-01-11 02:21:59.538118: I tensorflow/core/util/port.cc:110] oneDNN custom operations are on. You may see slightly different numerical results due to floating-point round-off errors from different computation orders. To turn them off, set the environment variable `TF_ENABLE_ONEDNN_OPTS=0`.
2024-01-11 02:21:59.575984: I tensorflow/core/platform/cpu_feature_guard.cc:182] This TensorFlow binary is optimized to use available CPU instructions in performance-critical operations.
To enable the following instructions: AVX2 AVX512F AVX512_VNNI FMA, in other operations, rebuild TensorFlow with the appropriate compiler flags.
2024-01-11 02:22:00.346338: W tensorflow/compiler/tf2tensorrt/utils/py_utils.cc:38] TF-TRT Warning: Could not find TensorRT
2024-01-11 02:22:02.022172: I tensorflow/compiler/xla/stream_executor/cuda/cuda_gpu_executor.cc:995] successful NUMA node read from SysFS had negative value (-1), but there must be at least one NUMA node, so returning NUMA node zero. See more at https://github.com/torvalds/linux/blob/v6.0/Documentation/ABI/testing/sysfs-bus-pci#L344-L355
2024-01-11 02:22:02.058030: I tensorflow/compiler/xla/stream_executor/cuda/cuda_gpu_executor.cc:995] successful NUMA node read from SysFS had negative value (-1), but there must be at least one NUMA node, so returning NUMA node zero. See more at https://github.com/torvalds/linux/blob/v6.0/Documentation/ABI/testing/sysfs-bus-pci#L344-L355
2024-01-11 02:22:02.061266: I tensorflow/compiler/xla/stream_executor/cuda/cuda_gpu_executor.cc:995] successful NUMA node read from SysFS had negative value (-1), but there must be at least one NUMA node, so returning NUMA node zero. See more at https://github.com/torvalds/linux/blob/v6.0/Documentation/ABI/testing/sysfs-bus-pci#L344-L355
2024-01-11 02:22:02.064864: I tensorflow/compiler/xla/stream_executor/cuda/cuda_gpu_executor.cc:995] successful NUMA node read from SysFS had negative value (-1), but there must be at least one NUMA node, so returning NUMA node zero. See more at https://github.com/torvalds/linux/blob/v6.0/Documentation/ABI/testing/sysfs-bus-pci#L344-L355
2024-01-11 02:22:02.067959: I tensorflow/compiler/xla/stream_executor/cuda/cuda_gpu_executor.cc:995] successful NUMA node read from SysFS had negative value (-1), but there must be at least one NUMA node, so returning NUMA node zero. See more at https://github.com/torvalds/linux/blob/v6.0/Documentation/ABI/testing/sysfs-bus-pci#L344-L355
2024-01-11 02:22:02.070999: I tensorflow/compiler/xla/stream_executor/cuda/cuda_gpu_executor.cc:995] successful NUMA node read from SysFS had negative value (-1), but there must be at least one NUMA node, so returning NUMA node zero. See more at https://github.com/torvalds/linux/blob/v6.0/Documentation/ABI/testing/sysfs-bus-pci#L344-L355
2024-01-11 02:22:02.820643: I tensorflow/compiler/xla/stream_executor/cuda/cuda_gpu_executor.cc:995] successful NUMA node read from SysFS had negative value (-1), but there must be at least one NUMA node, so returning NUMA node zero. See more at https://github.com/torvalds/linux/blob/v6.0/Documentation/ABI/testing/sysfs-bus-pci#L344-L355
2024-01-11 02:22:02.821876: I tensorflow/compiler/xla/stream_executor/cuda/cuda_gpu_executor.cc:995] successful NUMA node read from SysFS had negative value (-1), but there must be at least one NUMA node, so returning NUMA node zero. See more at https://github.com/torvalds/linux/blob/v6.0/Documentation/ABI/testing/sysfs-bus-pci#L344-L355
2024-01-11 02:22:02.822898: I tensorflow/compiler/xla/stream_executor/cuda/cuda_gpu_executor.cc:995] successful NUMA node read from SysFS had negative value (-1), but there must be at least one NUMA node, so returning NUMA node zero. See more at https://github.com/torvalds/linux/blob/v6.0/Documentation/ABI/testing/sysfs-bus-pci#L344-L355
2024-01-11 02:22:02.823877: I tensorflow/core/common_runtime/gpu/gpu_device.cc:1639] Created device /job:localhost/replica:0/task:0/device:GPU:0 with 13623 MB memory:  -> device: 0, name: Tesla T4, pci bus id: 0000:00:1e.0, compute capability: 7.5
dataset_len 600
2024-01-11 02:22:03.323582: W tensorflow/core/framework/op_kernel.cc:1828] OP_REQUIRES failed at strided_slice_op.cc:117 : INVALID_ARGUMENT: slice index 1 of dimension 0 out of bounds.
2024-01-11 02:22:03.323637: W tensorflow/core/framework/op_kernel.cc:1828] OP_REQUIRES failed at strided_slice_op.cc:117 : INVALID_ARGUMENT: Expected begin, end, and strides to be 1D equal size tensors, but got shapes [3], [1], and [1] instead.
2024-01-11 02:22:03.324399: W tensorflow/core/framework/op_kernel.cc:1828] OP_REQUIRES failed at strided_slice_op.cc:117 : INVALID_ARGUMENT: slice index 1 of dimension 0 out of bounds.
2024-01-11 02:22:03.324531: W tensorflow/core/framework/op_kernel.cc:1828] OP_REQUIRES failed at strided_slice_op.cc:117 : INVALID_ARGUMENT: Expected begin, end, and strides to be 1D equal size tensors, but got shapes [3], [1], and [1] instead.
2024-01-11 02:22:03.329642: W tensorflow/core/framework/op_kernel.cc:1816] UNKNOWN: InvalidArgumentError: {{function_node __wrapped__StridedSlice_device_/job:localhost/replica:0/task:0/device:GPU:0}} slice index 1 of dimension 0 out of bounds. [Op:StridedSlice] name: strided_slice/
Traceback (most recent call last):

  File "/home/ubuntu/miniconda3/envs/automl/lib/python3.9/site-packages/tensorflow/python/ops/script_ops.py", line 268, in __call__
    ret = func(*args)

  File "/home/ubuntu/miniconda3/envs/automl/lib/python3.9/site-packages/tensorflow/python/autograph/impl/api.py", line 643, in wrapper
    return func(*args, **kwargs)

  File "/home/ubuntu/automltraining/replicate_error.py", line 27, in prepare_input
    split_shape = tf.shape(image)[1] // 2

  File "/home/ubuntu/miniconda3/envs/automl/lib/python3.9/site-packages/tensorflow/python/util/traceback_utils.py", line 153, in error_handler
    raise e.with_traceback(filtered_tb) from None

  File "/home/ubuntu/miniconda3/envs/automl/lib/python3.9/site-packages/tensorflow/python/framework/ops.py", line 6656, in raise_from_not_ok_status
    raise core._status_to_exception(e) from None  # pylint: disable=protected-access

tensorflow.python.framework.errors_impl.InvalidArgumentError: {{function_node __wrapped__StridedSlice_device_/job:localhost/replica:0/task:0/device:GPU:0}} slice index 1 of dimension 0 out of bounds. [Op:StridedSlice] name: strided_slice/


2024-01-11 02:22:03.330620: W tensorflow/core/framework/op_kernel.cc:1816] UNKNOWN: InvalidArgumentError: {{function_node __wrapped__StridedSlice_device_/job:localhost/replica:0/task:0/device:GPU:0}} Expected begin, end, and strides to be 1D equal size tensors, but got shapes [3], [1], and [1] instead. [Op:StridedSlice] name: strided_slice/
Traceback (most recent call last):

  File "/home/ubuntu/miniconda3/envs/automl/lib/python3.9/site-packages/tensorflow/python/ops/script_ops.py", line 268, in __call__
    ret = func(*args)

  File "/home/ubuntu/miniconda3/envs/automl/lib/python3.9/site-packages/tensorflow/python/autograph/impl/api.py", line 643, in wrapper
    return func(*args, **kwargs)

  File "/home/ubuntu/automltraining/replicate_error.py", line 27, in prepare_input
    split_shape = tf.shape(image)[1] // 2

  File "/home/ubuntu/miniconda3/envs/automl/lib/python3.9/site-packages/tensorflow/python/util/traceback_utils.py", line 153, in error_handler
    raise e.with_traceback(filtered_tb) from None

  File "/home/ubuntu/miniconda3/envs/automl/lib/python3.9/site-packages/tensorflow/python/framework/ops.py", line 6656, in raise_from_not_ok_status
    raise core._status_to_exception(e) from None  # pylint: disable=protected-access

tensorflow.python.framework.errors_impl.InvalidArgumentError: {{function_node __wrapped__StridedSlice_device_/job:localhost/replica:0/task:0/device:GPU:0}} Expected begin, end, and strides to be 1D equal size tensors, but got shapes [3], [1], and [1] instead. [Op:StridedSlice] name: strided_slice/


2024-01-11 02:22:03.333474: W tensorflow/core/framework/op_kernel.cc:1816] UNKNOWN: InvalidArgumentError: {{function_node __wrapped__StridedSlice_device_/job:localhost/replica:0/task:0/device:GPU:0}} Expected begin, end, and strides to be 1D equal size tensors, but got shapes [3], [1], and [1] instead. [Op:StridedSlice] name: strided_slice/
Traceback (most recent call last):

  File "/home/ubuntu/miniconda3/envs/automl/lib/python3.9/site-packages/tensorflow/python/ops/script_ops.py", line 268, in __call__
    ret = func(*args)

  File "/home/ubuntu/miniconda3/envs/automl/lib/python3.9/site-packages/tensorflow/python/autograph/impl/api.py", line 643, in wrapper
    return func(*args, **kwargs)

  File "/home/ubuntu/automltraining/replicate_error.py", line 27, in prepare_input
    split_shape = tf.shape(image)[1] // 2

  File "/home/ubuntu/miniconda3/envs/automl/lib/python3.9/site-packages/tensorflow/python/util/traceback_utils.py", line 153, in error_handler
    raise e.with_traceback(filtered_tb) from None

  File "/home/ubuntu/miniconda3/envs/automl/lib/python3.9/site-packages/tensorflow/python/framework/ops.py", line 6656, in raise_from_not_ok_status
    raise core._status_to_exception(e) from None  # pylint: disable=protected-access

tensorflow.python.framework.errors_impl.InvalidArgumentError: {{function_node __wrapped__StridedSlice_device_/job:localhost/replica:0/task:0/device:GPU:0}} Expected begin, end, and strides to be 1D equal size tensors, but got shapes [3], [1], and [1] instead. [Op:StridedSlice] name: strided_slice/


2024-01-11 02:22:03.337583: W tensorflow/core/framework/op_kernel.cc:1816] UNKNOWN: InvalidArgumentError: {{function_node __wrapped__StridedSlice_device_/job:localhost/replica:0/task:0/device:GPU:0}} slice index 1 of dimension 0 out of bounds. [Op:StridedSlice] name: strided_slice/
Traceback (most recent call last):

  File "/home/ubuntu/miniconda3/envs/automl/lib/python3.9/site-packages/tensorflow/python/ops/script_ops.py", line 268, in __call__
    ret = func(*args)

  File "/home/ubuntu/miniconda3/envs/automl/lib/python3.9/site-packages/tensorflow/python/autograph/impl/api.py", line 643, in wrapper
    return func(*args, **kwargs)

  File "/home/ubuntu/automltraining/replicate_error.py", line 27, in prepare_input
    split_shape = tf.shape(image)[1] // 2

  File "/home/ubuntu/miniconda3/envs/automl/lib/python3.9/site-packages/tensorflow/python/util/traceback_utils.py", line 153, in error_handler
    raise e.with_traceback(filtered_tb) from None

  File "/home/ubuntu/miniconda3/envs/automl/lib/python3.9/site-packages/tensorflow/python/framework/ops.py", line 6656, in raise_from_not_ok_status
    raise core._status_to_exception(e) from None  # pylint: disable=protected-access

tensorflow.python.framework.errors_impl.InvalidArgumentError: {{function_node __wrapped__StridedSlice_device_/job:localhost/replica:0/task:0/device:GPU:0}} slice index 1 of dimension 0 out of bounds. [Op:StridedSlice] name: strided_slice/


2024-01-11 02:22:03.337628: W tensorflow/core/framework/op_kernel.cc:1816] UNKNOWN: InvalidArgumentError: {{function_node __wrapped__StridedSlice_device_/job:localhost/replica:0/task:0/device:GPU:0}} Expected begin, end, and strides to be 1D equal size tensors, but got shapes [3], [1], and [1] instead. [Op:StridedSlice] name: strided_slice/
Traceback (most recent call last):

  File "/home/ubuntu/miniconda3/envs/automl/lib/python3.9/site-packages/tensorflow/python/ops/script_ops.py", line 268, in __call__
    ret = func(*args)

  File "/home/ubuntu/miniconda3/envs/automl/lib/python3.9/site-packages/tensorflow/python/autograph/impl/api.py", line 643, in wrapper
    return func(*args, **kwargs)

  File "/home/ubuntu/automltraining/replicate_error.py", line 27, in prepare_input
    split_shape = tf.shape(image)[1] // 2

  File "/home/ubuntu/miniconda3/envs/automl/lib/python3.9/site-packages/tensorflow/python/util/traceback_utils.py", line 153, in error_handler
    raise e.with_traceback(filtered_tb) from None

  File "/home/ubuntu/miniconda3/envs/automl/lib/python3.9/site-packages/tensorflow/python/framework/ops.py", line 6656, in raise_from_not_ok_status
    raise core._status_to_exception(e) from None  # pylint: disable=protected-access

tensorflow.python.framework.errors_impl.InvalidArgumentError: {{function_node __wrapped__StridedSlice_device_/job:localhost/replica:0/task:0/device:GPU:0}} Expected begin, end, and strides to be 1D equal size tensors, but got shapes [3], [1], and [1] instead. [Op:StridedSlice] name: strided_slice/


Traceback (most recent call last):
  File "/home/ubuntu/automltraining/replicate_error.py", line 69, in <module>
    for x in dataset:
  File "/home/ubuntu/miniconda3/envs/automl/lib/python3.9/site-packages/tensorflow/python/data/ops/iterator_ops.py", line 814, in __next__
    return self._next_internal()
  File "/home/ubuntu/miniconda3/envs/automl/lib/python3.9/site-packages/tensorflow/python/data/ops/iterator_ops.py", line 777, in _next_internal
    ret = gen_dataset_ops.iterator_get_next(
  File "/home/ubuntu/miniconda3/envs/automl/lib/python3.9/site-packages/tensorflow/python/ops/gen_dataset_ops.py", line 3028, in iterator_get_next
    _ops.raise_from_not_ok_status(e, name)
  File "/home/ubuntu/miniconda3/envs/automl/lib/python3.9/site-packages/tensorflow/python/framework/ops.py", line 6656, in raise_from_not_ok_status
    raise core._status_to_exception(e) from None  # pylint: disable=protected-access
tensorflow.python.framework.errors_impl.UnknownError: {{function_node __wrapped__IteratorGetNext_output_types_2_device_/job:localhost/replica:0/task:0/device:CPU:0}} InvalidArgumentError: {{function_node __wrapped__StridedSlice_device_/job:localhost/replica:0/task:0/device:GPU:0}} slice index 1 of dimension 0 out of bounds. [Op:StridedSlice] name: strided_slice/
Traceback (most recent call last):

  File "/home/ubuntu/miniconda3/envs/automl/lib/python3.9/site-packages/tensorflow/python/ops/script_ops.py", line 268, in __call__
    ret = func(*args)

  File "/home/ubuntu/miniconda3/envs/automl/lib/python3.9/site-packages/tensorflow/python/autograph/impl/api.py", line 643, in wrapper
    return func(*args, **kwargs)

  File "/home/ubuntu/automltraining/replicate_error.py", line 27, in prepare_input
    split_shape = tf.shape(image)[1] // 2

  File "/home/ubuntu/miniconda3/envs/automl/lib/python3.9/site-packages/tensorflow/python/util/traceback_utils.py", line 153, in error_handler
    raise e.with_traceback(filtered_tb) from None

  File "/home/ubuntu/miniconda3/envs/automl/lib/python3.9/site-packages/tensorflow/python/framework/ops.py", line 6656, in raise_from_not_ok_status
    raise core._status_to_exception(e) from None  # pylint: disable=protected-access

tensorflow.python.framework.errors_impl.InvalidArgumentError: {{function_node __wrapped__StridedSlice_device_/job:localhost/replica:0/task:0/device:GPU:0}} slice index 1 of dimension 0 out of bounds. [Op:StridedSlice] name: strided_slice/


         [[{{node numpy_function_1}}]] [Op:IteratorGetNext] name:
```
