Repository: tensorflow
Issue: Tensorflow numpy_function causes errors when using tf.shape
Labels: stat:awaiting tensorflower,type:bug,comp:ops,TF 2.13

### 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:
```