Repository: tensorflow
Issue: ExtensionType to be considered a complete nested structure
Labels: stat:awaiting tensorflower,type:feature,comp:apis

### Issue type

Feature Request

### Have you reproduced the bug with TensorFlow Nightly?

Yes

### Source

source

### TensorFlow version

2.13  (& 2.15 nightly)

### Custom code

No

### OS platform and distribution

_No response_

### Mobile device

_No response_

### Python version

3.10

### Bazel version

_No response_

### GCC/compiler version

_No response_

### CUDA/cuDNN version

_No response_

### GPU model and memory

_No response_

### Current behavior?

Tensorflow's current experimental [ExtensionType](https://www.tensorflow.org/api_docs/python/tf/experimental/ExtensionType) is a very useful data structure to replace python data structures when the modelled tensors are in fact more complex structures of tensors of varying shapes and types. One would intuitively think that this class would click with `tf.nest`, but unfortunately they are not among the valid nested structures:

```
collections.abc.Sequence (except string and bytes). This includes list, tuple, and namedtuple.
collections.abc.Mapping (with sortable keys). This includes dict and collections.OrderedDict.
collections.abc.MappingView (with sortable keys).
```

as taken from [Module: tf.nest](https://www.tensorflow.org/api_docs/python/tf/nest).

Moreover, `tf.nest` is not registered as an api for [@tf.experimental.dispatch_for_api(tf_api)](https://www.tensorflow.org/api_docs/python/tf/experimental/dispatch_for_api), which would be helpful for redifining the behaviour of nesting inside an `ExtensionType`.

Currently, the only solutions that I have found were either to:

1. Use multiple inheritance and also inherit from one of the valid nested structures in order to allow `tf.nest` to work as an api. This also requires implementing all the abstract methods of the superclass which may not make sense all the time.
2. Rewrite all the functionalities that use nesting as methods inside the class implementing `ExtensionType`, which is not great when you are seeking code reusability.
3. Have a method to change the type to something that can be considered a structure and have that called whenever `tf.nest`ing is used, which as well requires some changes everywhere in the code base. 


### Standalone code to reproduce the issue

```python
# Not inheriting from Sequence

import tensorflow as tf

class MyComplexData(tf.experimental.ExtensionType):
    x: tf.Tensor
    y: tf.Tensor

    def __init__(self, x, y):
        self.x = x
        self.y = y

def gather_fn(tensor):
    indices = tf.constant([[0], [1]])
    return tf.gather_nd(tensor, indices)

data = MyComplexData(tf.constant([1, 2, 3]), tf.constant([[1.0, 2.0], [3.0, 4.0]]))

try:
    mapped_data = tf.nest.map_structure(gather_fn, data)
    print(mapped_data)
except Exception as e:
    print(f"Failed to apply map_structure with gather_fn on MyComplexData: {e}")
```
```python
# Inheriting from Sequence

import tensorflow as tf
from collections.abc import Sequence

class MyComplexDataSequence(tf.experimental.ExtensionType, Sequence):
    x: tf.Tensor
    y: tf.Tensor

    def __init__(self, x, y):
        self.x = x
        self.y = y

    def __getitem__(self, index):
        return [self.x, self.y][index]
    
    def __len__(self):
        return 2

def gather_fn(tensor):
    indices = tf.constant([[0], [1]])
    return tf.gather_nd(tensor, indices)

data_seq = MyComplexDataSequence(tf.constant([1, 2, 3]), tf.constant([[1.0, 2.0], [3.0, 4.0]]))

try:
    mapped_data_seq = tf.nest.map_structure(gather_fn, data_seq)
    print(mapped_data_seq)
except Exception as e:
    print(f"Failed to apply map_structure with gather_fn on MyComplexDataSequence: {e}")
```


### Relevant log output

```shell
# Not inheriting from Sequence

Failed to apply map_structure with gather_fn on MyComplexData: Attempt to convert a value (MyComplexData(x=<tf.Tensor: shape=(3,), dtype=int32, numpy=array([1, 2, 3], dtype=int32)>, y=<tf.Tensor: shape=(2, 2), dtype=float32, numpy=
array([[1., 2.],
       [3., 4.]], dtype=float32)>)) with an unsupported type (<class '__main__.MyComplexData'>) to a Tensor.
```

```shell
# Inheriting from Sequence

MyComplexDataSequence(x=<tf.Tensor: shape=(2,), dtype=int32, numpy=array([1, 2], dtype=int32)>, y=<tf.Tensor: shape=(2, 2), dtype=float32, numpy=
array([[1., 2.],
       [3., 4.]], dtype=float32)>)
```