Repository: tensorflow
Issue: Dataset.ragged_batch does not produce correct specs with tf.py_function and tf.numpy_function
Labels: stat:awaiting tensorflower,type:bug,comp:data,TF 2.12

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

Bug

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

Yes

### Source

source

### Tensorflow Version

2.12

### Custom Code

Yes

### OS Platform and Distribution

docker container nvcr.io/nvidia/tensorflow:23.04-tf2-py3 on Ubuntu 22.04.2 LTS host

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

I'm trying to train an object detection model where images may have a different number of bounding boxes. Also I want to add some augmentations, and since tf does not support augmentations of bounding boxes I choose albumentations to do the job. I can't use albumentations' augmentations directly, so I need to use either `tf.py_function` or `tf.numpy_function`. I used `Dataset.ragged_batch` instead of `Dataset.batch` (because the dimension of bbox tensor may vary), but it did not provide me the correct `element_spec` and I was unable to make it work.

These are three scenarios that should help to understand the issue:

### Scenario 1:
I don't use any augmentations, `ragged_batch` returns the correct element spec, but I really need those augmentations
```
(TensorSpec(shape=(None, 512, 512, 3), dtype=tf.float32, name=None),
 {'classes': RaggedTensorSpec(TensorShape([None, None]), tf.float32, 1, tf.int64),
  'boxes': RaggedTensorSpec(TensorShape([None, None, 4]), tf.float64, 1, tf.int64)})
```
### Scenario 2:
I use `tf.numpy_function` fo perform the augmentations. The spec is incorrect, I can't batch the items
```
(TensorSpec(shape=<unknown>, dtype=tf.float32, name=None),
 {'classes': TensorSpec(shape=<unknown>, dtype=tf.float32, name=None),
  'boxes': TensorSpec(shape=<unknown>, dtype=tf.float32, name=None)})
```
### Scenario 3
I use `tf.py_function`, provide something, that looks like correct spec to `Tout` param:
```
Tout=[
    tf.TensorSpec(shape=[None, 512, 512, 3], dtype=tf.float32), 
    tf.RaggedTensorSpec(shape=[None, 4], dtype=tf.float32), 
    tf.RaggedTensorSpec(shape=[None, None], dtype=tf.float32),
],
```
but spec for image still does not look good
```
(TensorSpec(shape=<unknown>, dtype=tf.float32, name=None),
 {'classes': RaggedTensorSpec(TensorShape([None, None, None]), tf.float32, 2, tf.int64),
  'boxes': RaggedTensorSpec(TensorShape([None, None, 4]), tf.float32, 2, tf.int64)})
```
and I had to add an extra dimension for labels spec in order to convert it to `RaggedTensor` (which is probably not good as well). Model refuses to be trained because of incorrect image dimensions

The (non)working code is here - [colab link](https://colab.research.google.com/drive/148i78QRnF98guvx1Y0gZXtBVrv-day42?usp=sharing)

### Standalone code to reproduce the issue

```shell
import keras_cv
import numpy as np
import tensorflow as tf
import albumentations as A

def generate_random_data():
    while True:
        image = np.random.randint(0, 256, size=(512, 512, 3), dtype=np.uint8)
        num_bboxes = np.random.randint(1, 200)
        bboxes = []
        labels = []
        for _ in range(num_bboxes):
            x_min, y_min, x_max, y_max = np.sort(np.random.uniform(0, 512, size=4) / 512)
            bbox = [x_min, y_min, x_max, y_max]
            label = np.random.choice([0, 1])
            bboxes.append(bbox)
            labels.append(label)

        data = {
            'image': tf.convert_to_tensor(image),
            'bboxes': {
                'bbox': tf.convert_to_tensor(bboxes),
                'label': tf.convert_to_tensor(labels, dtype=tf.int64),
            }
        }
        yield data

# Create the random dataset
dataset = tf.data.Dataset.from_generator(generate_random_data, output_signature={
    'image': tf.TensorSpec(shape=(512, 512, 3), dtype=tf.uint8),
    'bboxes': {
        'bbox': tf.TensorSpec(shape=(None, 4), dtype=tf.float64),
        'label': tf.TensorSpec(shape=(None,), dtype=tf.int64),
    },
})


# Scenario 1
def preprocess_data_1(inputs):
    bounding_boxes = {
        "classes": tf.cast(inputs["bboxes"]["label"], dtype=tf.float32),
        "boxes": inputs["bboxes"]["bbox"],
    }
    return tf.image.convert_image_dtype(inputs['image'], tf.float32), bounding_boxes
ds_1 = dataset.map(preprocess_data_1).ragged_batch(2)

# Scenario 2
def transform_2(image, bboxes, labels):
    transforms = A.Compose(
        [
            A.Rotate(limit=40),
        ],
        bbox_params=A.BboxParams(
            format='albumentations',
            label_fields=['label'],
        )
    )

    transformed = transforms(
        image=image, 
        label=labels,
        bboxes=bboxes,
    )

    return transformed

def aug_fn_2(image, bboxes, labels):
    aug_data = transform_2(image, bboxes, labels)
    return (
        tf.image.convert_image_dtype(aug_data["image"], tf.float32), 
        tf.convert_to_tensor(aug_data["bboxes"], dtype=tf.float32), 
        tf.cast(aug_data["label"], tf.float32)
    )

def preprocess_data_2(inputs):
    bboxes = inputs['bboxes']['bbox']
    labels = inputs['bboxes']['label']
    aug_image, aug_bboxes, aug_labels = tf.numpy_function(
        func=aug_fn_2, 
        inp=[inputs["image"], bboxes, labels], 
        Tout=[tf.float32, tf.float32, tf.float32],
    )
    
    bounding_boxes = {
        "classes": aug_labels,
        "boxes":  aug_bboxes
    }

    return aug_image, bounding_boxes

ds_2 = dataset.map(preprocess_data_2).ragged_batch(2)
for item in ds_2:
    break

# Scenario 3
def transform_3(image, bboxes, labels):
    transforms = A.Compose(
        [
            A.Rotate(limit=40),
        ],
        bbox_params=A.BboxParams(
            format='albumentations',
            label_fields=['label'],
        )
    )

    transformed = transforms(
        image=image.numpy(), 
        label=labels.numpy(),
        bboxes=bboxes.numpy(),
    )

    return transformed

def aug_fn_3(image, bboxes, labels):
    aug_data = transform_3(image, bboxes, labels)

    return (
        tf.image.convert_image_dtype(aug_data["image"], tf.float32), 
        tf.RaggedTensor.from_tensor(tf.convert_to_tensor(aug_data["bboxes"], dtype=tf.float32)),
        tf.RaggedTensor.from_tensor(tf.cast([aug_data["label"]], tf.float32)),
    )

def preprocess_data_3(inputs):
    bboxes = inputs['bboxes']['bbox']
    labels = inputs['bboxes']['label']
    aug_image, aug_bboxes, aug_labels = tf.py_function(
        func=aug_fn_3, 
        inp=[inputs["image"], bboxes, labels], 
        Tout=[
            tf.TensorSpec(shape=[None, 512, 512, 3], dtype=tf.float32), 
            tf.RaggedTensorSpec(shape=[None, 4], dtype=tf.float32), 
            tf.RaggedTensorSpec(shape=[None, None], dtype=tf.float32),
        ],
    )
    
    bounding_boxes = {
        "classes": aug_labels,
        "boxes":  aug_bboxes,
    }

    return aug_image, bounding_boxes

ds_3 = dataset.map(preprocess_data_3).ragged_batch(2)
for batch in ds_3:
    break

model = keras_cv.models.RetinaNet.from_preset(
    "resnet50_imagenet",
    num_classes=2,
    bounding_box_format="rel_xyxy",
)

model.compile(
    classification_loss="focal",
    box_loss="smoothl1",
    optimizer=tf.optimizers.Adam(),
)

model.fit(
    ds_3.take(1),
    validation_data=ds_3.take(1),
    epochs=100,
)
```


### Relevant log output

_No response_</details>