rejection_resample loses track of ragged tensors
It appears this can be worked around like this:
```python
input = tf.ragged.constant([[1, 2, 3], [4, 5], [7, 8, 9]])
ds = tf.data.Dataset.from_tensor_slices(input)
original_spec = ds.element_spec
tf.random.set_seed(0)
ds = ds.rejection_resample(
    class_func=lambda t: 1,
    target_dist=(0.1, 0.9),
).map(lambda class_func_result, data: data)
ds = ds.batch(1).map(lambda elem: tf.map_fn(
      fn=lambda e: e,
      elems=elem,
      fn_output_signature=original_spec,
  )).unbatch()
ds = ds.batch(2)
ds.take(1).get_single_element()
```

But that's pretty annoying. IDK if there's a shorter way to express the same thing.

Also, it appears that the same problem applies generally to `Dataset.map`, i.e. this breaks in the same way:
```python
input = tf.ragged.constant([[1, 2, 3], [4, 5], [7, 8, 9]])
ds = tf.data.Dataset.from_tensor_slices(input)
ds = ds.map(lambda e: e)
ds = ds.batch(2)
ds.take(1).get_single_element()
``` Hi @chris-remedy ,

I have replicated the reported behaviour and attached [gist](https://colab.research.google.com/gist/SuryanarayanaY/6e2138c261058cb31dfae033f5a7061b/60583.ipynb) for reference. This needs to dig into more to confirm the root cause. Thanks!
<details><summary>Click to expand!</summary> 
 
 ### Issue Type

Bug

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

Yes

### Source

binary

### Tensorflow Version

2.14.0-dev20230512

### Custom Code

No

### OS Platform and Distribution

Ubuntu 20.04.6

### Mobile device

_No response_

### Python version

3.8.14

### Bazel version

_No response_

### GCC/Compiler version

_No response_

### CUDA/cuDNN version

_No response_

### GPU model and memory

_No response_

### Current Behaviour?

A `tf.data.Dataset` initialized from RaggedTensors normally will successfully batch into ragged batches. However after passing it through `rejection_resample`, it loses track of which input tensors were ragged, and so batching fails.

### Standalone code to reproduce the issue

```shell
import tensorflow as tf

input = tf.ragged.constant([[1,2,3], [4,5], [7,8,9]])
ds = tf.data.Dataset.from_tensor_slices(input)
tf.random.set_seed(0)
# Removing this line makes everything work fine
ds = ds.rejection_resample(
    class_func=lambda t: 1,
    target_dist=(0.1, 0.9),
)
ds = ds.batch(2)
ds.take(1).get_single_element()
```


### Relevant log output

```shell
tensorflow.python.framework.errors_impl.InvalidArgumentError: {{function_node __wrapped__DatasetToSingleElement_output_types_2_device_/job:localhost/replica:0/task:0/device:CPU:0}} Cannot batch tensors with different shapes in component 1. First element had shape [3] and element 1 had shape [2]. [Op:DatasetToSingleElement] name:
```
</details>