ExtensionType to be considered a complete nested structure
I volunteer to take this issue The error you're encountering is because tf.nest.map_structure requires all elements within the nested structure to be TensorFlow tensors, and MyComplexData is not considered a TensorFlow tensor by default. To make your code compatible with tf.nest.map_structure, you need to ensure that the elements within MyComplexData are TensorFlow tensors.

You can do this by adding a method to the MyComplexData class that converts its internal data to TensorFlow tensors. In the updated code, I've included a method called as_tensors that does this conversion for x and y. By calling data.as_tensors(), you can convert your custom data structure into a format compatible with TensorFlow operations. This should allow you to use tf.nest.map_structure with the gather_fn function without encountering the error. 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 as_tensors(self):
        return MyComplexData(tf.convert_to_tensor(self.x), tf.convert_to_tensor(self.y))

def gather_fn(tensor):
    indices = tf.constant([[0], [1]])
    return tf.gather_nd(tensor, indices)

data = MyComplexData([1, 2, 3], [[1.0, 2.0], [3.0, 4.0]])

# Convert the data to TensorFlow tensors
data = data.as_tensors()

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}")

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