The inferred shape of `tf.RaggedTensor.row_lengths(axis=2)` in Keras graph is incorrect for ragged tensor with uniform row lengths
Hi @foxik
Thank you for reporting this issue!

I was able to replicate the issue in Colab using Tensorflow 2.12 and tf-nightly(2.13.0.dev20230424). Please find the gists - [TF 2.12](https://colab.sandbox.google.com/gist/synandi/3307d76df72bd5181ed5aa4dd639b7ac/60400.ipynb) & [tf-nightly](https://colab.sandbox.google.com/gist/synandi/fef428141788b9021bb529750e6d7f1c/60400_nightly.ipynb). 
Thank you!
  Hi @foxik ,

The difference in results are due to 2 reasons.The code `inputs.row_lengths(axis=1)` outputs `KerasTensor` with `type_spec=TensorSpec` whereas inputs.row_lengths(axis=2) outputs` KerasTensor(type_spec=RaggedTensorSpec)`. Please note the difference here `TensorSpec` vs `RaggedTensorSpec`.

When we apply index/slicing operation [:] , internally both calls __getitem__ method and you can find from the colab gist attached in [comment](https://github.com/tensorflow/tensorflow/issues/60400#issuecomment-1520309736), for axis=1 tf.__operators__.getitem' method called and for axis=2 'tf.__operators__.ragged_getitem' called and both are different methods and hence different results.

```
print(inputs.row_lengths(axis=1))
#KerasTensor(type_spec=TensorSpec(shape=(None,), dtype=tf.int64, name=None), name='input.row_lengths/sub:0', description="created by layer 'input.row_lengths'")

print(inputs.row_lengths(axis=1)[:]) # OK, is the same as above
#KerasTensor(type_spec=TensorSpec(shape=(None,), dtype=tf.int64, name=None), name='tf.__operators__.getitem/strided_slice:0', description="created by layer `'tf.__operators__.getitem'")`

print(inputs.row_lengths(axis=2))
#KerasTensor(type_spec=RaggedTensorSpec(TensorShape([None, 64]), tf.int64, 1, tf.int64), description="created by layer 'input.row_lengths_2'")

print(inputs.row_lengths(axis=2)[:]) # Not the same as above as it calls different methods
#KerasTensor(type_spec=RaggedTensorSpec(TensorShape([1, 64]), tf.int64, 1, tf.int64), description="created by layer 'tf.__operators__.ragged_getitem'")
```
Here the axis and the type spec making the difference.

However, I have tested the same with actual data and observed same result with same specs and values.Please refer the attached [gist](https://colab.research.google.com/gist/SuryanarayanaY/0c013dfafc19261d3dab12d711d21293/60400_r1.ipynb) with the same operations performed on actual ragged tensor.Can you test also with some actual data and confirm if there is any difference. Thanks for your answer!

The problem is that
```python
inputs.row_lengths(axis=2).shape != inputs.row_lengths(axis=2)[:].shape
```
while it should be the same, because for any tensor (or ragged tensor) `a`, it should be the case that `a.shape == (a[:]).shape`.

I understand the difference between `inputs.row_lengths(axis=1)` and `inputs.row_lengths(axis=2)`; I just included it in the example to show that for `axis=1` it is indeed the case that
```python
inputs.row_lengths(axis=1).shape == inputs.row_lengths(axis=1)[:].shape
```

The reported problem concerns only static shape inference; during the graph evaluation, the shape is correct. That means that even if `inputs.row_lengths(axis=2)[:].shape[0]` is `1`, the `tf.shape(inputs.row_lengths(axis=2)[:])[0]` can be something different than 1 (in the gist it is 5), which is incorrect -- such incorrect static shape can cause problems during graph construction.

We encountered the problem in a OCR-like task -- given images of fixed height and dynamic width as ragged tensors, we wanted to extract the widths of the images. That can be achieved using for example `inputs.row_lengths(axis=2)[:, 0]` -- but the `.shape` of that is reported as `(1,)`, which is clearly incorrect.

So I believe this behavior is indeed a bug.
<details><summary>Click to expand!</summary> 
 
 ### Issue Type

Bug

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

Yes

### Source

binary

### Tensorflow Version

TF 2.12, TF nightly 2.13.0-dev20230420

### Custom Code

No

### OS Platform and Distribution

_No response_

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

Assume a Keras graph gets a ragged tensor with uniform row lengths in the axis 1, so for example
```python
inputs = tf.keras.layers.Input([64, None], ragged=True)
```

Then `inputs.row_lengths(axis=2)` is a `KerasTensor` with a correct spec `RaggedTensorSpec(TensorShape([None, 64]), ...)`.

However, when you index the first axis using a "full" slice, i.e.
```python
inputs.row_lengths(axis=2)[:]
```
you should get the same tensor -- but you get a `KerasTensor` with an incorrect spec `RaggedTensorSpec(TensorShape([1, 64]), ...)`.

### Standalone code to reproduce the issue

A Colab notebook reproducing the issue both in TF 2.12.0 and in TF nightly 2.13.0-dev20230420 can be found at https://colab.research.google.com/drive/1RKvNdB_81yKZkfzDpefIPJU2pgqZuYUY?usp=sharing

The full source also follows:

```python
inputs = tf.keras.layers.Input([64, None], ragged=True)
print(inputs)

print(inputs.row_lengths(axis=1))
print(inputs.row_lengths(axis=1)[:]) # OK, is the same as above

print(inputs.row_lengths(axis=2))
print(inputs.row_lengths(axis=2)[:]) # Problem, should be the same as above
```


### Relevant log output

```python
KerasTensor(type_spec=RaggedTensorSpec(TensorShape([None, 64, None]), tf.float32, 2, tf.int64), name='input_1', description="created by layer 'input_1'")
KerasTensor(type_spec=TensorSpec(shape=(None,), dtype=tf.int64, name=None), name='input.row_lengths/sub:0', description="created by layer 'input.row_lengths'")
KerasTensor(type_spec=TensorSpec(shape=(None,), dtype=tf.int64, name=None), name='tf.__operators__.getitem/strided_slice:0', description="created by layer 'tf.__operators__.getitem'")
KerasTensor(type_spec=RaggedTensorSpec(TensorShape([None, 64]), tf.int64, 1, tf.int64), description="created by layer 'input.row_lengths_2'")
KerasTensor(type_spec=RaggedTensorSpec(TensorShape([1, 64]), tf.int64, 1, tf.int64), description="created by layer 'tf.__operators__.ragged_getitem'")
```
</details>