Nested ExtensionTypes failing with new __tf_flatten__ / __tf_unflatten__ api
Hi @areiner222 ,

I have replicated the reported issue and attached gist for [reference](https://colab.sandbox.google.com/gist/SuryanarayanaY/a31e40d2e2cca8f6c71d5a704eb17240/62082.ipynb#scrollTo=Xg_ZS9IIy1wn). We need to check this.

Thanks!
### Issue type

Bug

### Have you reproduced the bug with TensorFlow Nightly?

No

### Source

source

### TensorFlow version

tf 2.14

### Custom code

Yes

### OS platform and distribution

Apple Silicon M1

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

Nested ExtensionTypes failing with new __tf_flatten__, __tf_unflatten__ api

### Standalone code to reproduce the issue

```shell

import dataclasses
import tensorflow as tf

@dataclasses.dataclass
class MaskedTensor(tf.experimental.ExtensionType):
    mask: bool
    value: tf.Tensor

    def __tf_flatten__(self):
        metadata = (self.mask,)  # static config.
        components = (self.value,)  # dynamic values.
        return metadata, components

    @classmethod
    def __tf_unflatten__(cls, metadata, components):
        return cls(*metadata, *components)
        

@dataclasses.dataclass
class MaskedTensorComp(tf.experimental.ExtensionType):
    mask: bool
    value: tf.Tensor
    mt: MaskedTensor

    def __tf_flatten__(self):
        metadata = (self.mask)  # static config.
        components = (self.value, self.mt)  # dynamic values.
        return metadata, components

    @classmethod
    def __tf_unflatten__(cls, metadata, components):
        print('Unflattening MaskedTensorComposite', components, metadata)
        return cls(*metadata, *components)

mt0 = MaskedTensor(True, tf.constant(3.))
mt = MaskedTensorComp(False, tf.constant(99.), mt=mt0)

mt_flat = tf.nest.flatten(mt)
mt_recon = tf.nest.pack_sequence_as(mt, mt_flat)
```
```


### Relevant log output

```shell
File ~/mambaforge/envs/_/lib/python3.10/site-packages/tensorflow/python/util/nest.py:538, in pack_sequence_as(structure, flat_sequence, expand_composites)
    424 @tf_export("nest.pack_sequence_as")
    425 def pack_sequence_as(structure, flat_sequence, expand_composites=False):
    426   """Returns a given flattened sequence packed into a given structure.
    427 
    428   Refer to [tf.nest](https://www.tensorflow.org/api_docs/python/tf/nest)
   (...)
    536     TypeError: `structure` is or contains a dict with non-sortable keys.
    537   """
--> 538   return nest_util.pack_sequence_as(
    539       nest_util.Modality.CORE, structure, flat_sequence, expand_composites
    540   )

File ~/mambaforge/envs//lib/python3.10/site-packages/tensorflow/python/util/nest_util.py:958, in pack_sequence_as(modality, structure, flat_sequence, expand_composites, sequence_fn)
    835 """Returns a given flattened sequence packed into a given structure.
    836 
    837 - For Modality.CORE: Refer to
   (...)
    955   non-sortable keys.
    956 """
    957 if modality == Modality.CORE:
--> 958   return _tf_core_pack_sequence_as(
    959       structure, flat_sequence, expand_composites, sequence_fn
    960   )
    961 elif modality == Modality.DATA:
    962   return _tf_data_pack_sequence_as(structure, flat_sequence)

File ~/mambaforge/envs//lib/python3.10/site-packages/tensorflow/python/util/nest_util.py:1022, in _tf_core_pack_sequence_as(structure, flat_sequence, expand_composites, sequence_fn)
   1015   if len(flat_structure) != len(flat_sequence):
   1016     # pylint: disable=raise-missing-from
   1017     raise ValueError(
   1018         "Could not pack sequence. Structure had %d atoms, but "
   1019         "flat_sequence had %d items.  Structure: %s, flat_sequence: %s."
   1020         % (len(flat_structure), len(flat_sequence), structure, flat_sequence)
   1021     )
-> 1022 return sequence_fn(structure, packed)

File ~/mambaforge/envs//lib/python3.10/site-packages/tensorflow/python/util/nest_util.py:336, in sequence_like(instance, args)
    334   assert len(args) == 1
    335   spec = instance._type_spec  # pylint: disable=protected-access
--> 336   return spec._from_components(args[0])  # pylint: disable=protected-access
    337 elif _is_type_spec(instance):
    338   # Pack a CompositeTensor's components according to a TypeSpec.
    339   assert len(args) == 1

File ~/mambaforge/envs//lib/python3.10/site-packages/tensorflow/python/framework/extension_type.py:482, in ExtensionTypeSpec._from_components(self, components)
    476 if list(components_iter):
    477   raise ValueError(
    478       'Cannot build an ExtensionType instance from components '
    479       'because more components are provided than the number expected '
    480       'by the type spec.'
    481   )
--> 482 value_tuple = nest.pack_sequence_as(spec_tuple, flat)
    483 fields = dict(zip(self.__dict__.keys(), value_tuple))
    485 # Build the new value.  Bypass the constructor (__init__), in case the user
    486 # who defined the ExtensionType used a custom constructor.

File ~/mambaforge/envs//lib/python3.10/site-packages/tensorflow/python/util/nest.py:538, in pack_sequence_as(structure, flat_sequence, expand_composites)
    424 @tf_export("nest.pack_sequence_as")
    425 def pack_sequence_as(structure, flat_sequence, expand_composites=False):
    426   """Returns a given flattened sequence packed into a given structure.
    427 
    428   Refer to [tf.nest](https://www.tensorflow.org/api_docs/python/tf/nest)
   (...)
    536     TypeError: `structure` is or contains a dict with non-sortable keys.
    537   """
--> 538   return nest_util.pack_sequence_as(
    539       nest_util.Modality.CORE, structure, flat_sequence, expand_composites
    540   )

File ~/mambaforge/envs//lib/python3.10/site-packages/tensorflow/python/util/nest_util.py:958, in pack_sequence_as(modality, structure, flat_sequence, expand_composites, sequence_fn)
    835 """Returns a given flattened sequence packed into a given structure.
    836 
    837 - For Modality.CORE: Refer to
   (...)
    955   non-sortable keys.
    956 """
    957 if modality == Modality.CORE:
--> 958   return _tf_core_pack_sequence_as(
    959       structure, flat_sequence, expand_composites, sequence_fn
    960   )
    961 elif modality == Modality.DATA:
    962   return _tf_data_pack_sequence_as(structure, flat_sequence)

File ~/mambaforge/envs//lib/python3.10/site-packages/tensorflow/python/util/nest_util.py:1006, in _tf_core_pack_sequence_as(structure, flat_sequence, expand_composites, sequence_fn)
   1003   return flat_sequence[0]
   1005 try:
-> 1006   final_index, packed = _tf_core_packed_nest_with_indices(
   1007       structure, flat_sequence, 0, is_nested_fn, sequence_fn
   1008   )
   1009   if final_index < len(flat_sequence):
   1010     raise IndexError

File ~/mambaforge/envs//lib/python3.10/site-packages/tensorflow/python/util/nest_util.py:679, in _tf_core_packed_nest_with_indices(structure, flat, index, is_nested_fn, sequence_fn)
    675 if is_nested_fn(s):
    676   new_index, child = _tf_core_packed_nest_with_indices(
    677       s, flat, index, is_nested_fn, sequence_fn
    678   )
--> 679   packed.append(sequence_fn(s, child))
    680   index = new_index
    681 else:

File ~/mambaforge/envs//lib/python3.10/site-packages/tensorflow/python/util/nest_util.py:336, in sequence_like(instance, args)
    334   assert len(args) == 1
    335   spec = instance._type_spec  # pylint: disable=protected-access
--> 336   return spec._from_components(args[0])  # pylint: disable=protected-access
    337 elif _is_type_spec(instance):
    338   # Pack a CompositeTensor's components according to a TypeSpec.
    339   assert len(args) == 1

File ~/mambaforge/envs//lib/python3.10/site-packages/tensorflow/python/framework/extension_type.py:472, in ExtensionTypeSpec._from_components(self, components)
    470 spec_tuple = tuple(self.__dict__.values())
    471 components_iter = iter(components)
--> 472 flat = [
    473     next(components_iter) if isinstance(x, type_spec.TypeSpec) else x
    474     for x in nest.flatten(spec_tuple)
    475 ]
    476 if list(components_iter):
    477   raise ValueError(
    478       'Cannot build an ExtensionType instance from components '
    479       'because more components are provided than the number expected '
    480       'by the type spec.'
    481   )

File ~/mambaforge/envs/
File ~/mambaforge/envs//lib/python3.10/site-packages/tensorflow/python/util/nest.py:538, in pack_sequence_as(structure, flat_sequence, expand_composites)
    424 @tf_export("nest.pack_sequence_as")
    425 def pack_sequence_as(structure, flat_sequence, expand_composites=False):
    426   """Returns a given flattened sequence packed into a given structure.
    427 
    428   Refer to [tf.nest](https://www.tensorflow.org/api_docs/python/tf/nest)
   (...)
    536     TypeError: `structure` is or contains a dict with non-sortable keys.
    537   """
--> 538   return nest_util.pack_sequence_as(
    539       nest_util.Modality.CORE, structure, flat_sequence, expand_composites
    540   )

File ~/mambaforge/envs//lib/python3.10/site-packages/tensorflow/python/util/nest_util.py:958, in pack_sequence_as(modality, structure, flat_sequence, expand_composites, sequence_fn)
    835 """Returns a given flattened sequence packed into a given structure.
    836 
    837 - For Modality.CORE: Refer to
   (...)
    955   non-sortable keys.
    956 """
    957 if modality == Modality.CORE:
--> 958   return _tf_core_pack_sequence_as(
    959       structure, flat_sequence, expand_composites, sequence_fn
    960   )
    961 elif modality == Modality.DATA:
    962   return _tf_data_pack_sequence_as(structure, flat_sequence)

File ~/mambaforge/envs//lib/python3.10/site-packages/tensorflow/python/util/nest_util.py:1022, in _tf_core_pack_sequence_as(structure, flat_sequence, expand_composites, sequence_fn)
   1015   if len(flat_structure) != len(flat_sequence):
   1016     # pylint: disable=raise-missing-from
   1017     raise ValueError(
   1018         "Could not pack sequence. Structure had %d atoms, but "
   1019         "flat_sequence had %d items.  Structure: %s, flat_sequence: %s."
   1020         % (len(flat_structure), len(flat_sequence), structure, flat_sequence)
   1021     )
-> 1022 return sequence_fn(structure, packed)

File ~/mambaforge/envs//lib/python3.10/site-packages/tensorflow/python/util/nest_util.py:336, in sequence_like(instance, args)
    334   assert len(args) == 1
    335   spec = instance._type_spec  # pylint: disable=protected-access
--> 336   return spec._from_components(args[0])  # pylint: disable=protected-access
    337 elif _is_type_spec(instance):
    338   # Pack a CompositeTensor's components according to a TypeSpec.
    339   assert len(args) == 1

File ~/mambaforge/envs//lib/python3.10/site-packages/tensorflow/python/framework/extension_type.py:482, in ExtensionTypeSpec._from_components(self, components)
    476 if list(components_iter):
    477   raise ValueError(
    478       'Cannot build an ExtensionType instance from components '
    479       'because more components are provided than the number expected '
    480       'by the type spec.'
    481   )
--> 482 value_tuple = nest.pack_sequence_as(spec_tuple, flat)
    483 fields = dict(zip(self.__dict__.keys(), value_tuple))
    485 # Build the new value.  Bypass the constructor (__init__), in case the user
    486 # who defined the ExtensionType used a custom constructor.

File ~/mambaforge/envs//lib/python3.10/site-packages/tensorflow/python/util/nest.py:538, in pack_sequence_as(structure, flat_sequence, expand_composites)
    424 @tf_export("nest.pack_sequence_as")
    425 def pack_sequence_as(structure, flat_sequence, expand_composites=False):
    426   """Returns a given flattened sequence packed into a given structure.
    427 
    428   Refer to [tf.nest](https://www.tensorflow.org/api_docs/python/tf/nest)
   (...)
    536     TypeError: `structure` is or contains a dict with non-sortable keys.
    537   """
--> 538   return nest_util.pack_sequence_as(
    539       nest_util.Modality.CORE, structure, flat_sequence, expand_composites
    540   )

File ~/mambaforge/envs//lib/python3.10/site-packages/tensorflow/python/util/nest_util.py:958, in pack_sequence_as(modality, structure, flat_sequence, expand_composites, sequence_fn)
    835 """Returns a given flattened sequence packed into a given structure.
    836 
    837 - For Modality.CORE: Refer to
   (...)
    955   non-sortable keys.
    956 """
    957 if modality == Modality.CORE:
--> 958   return _tf_core_pack_sequence_as(
    959       structure, flat_sequence, expand_composites, sequence_fn
    960   )
    961 elif modality == Modality.DATA:
    962   return _tf_data_pack_sequence_as(structure, flat_sequence)

File ~/mambaforge/envs//lib/python3.10/site-packages/tensorflow/python/util/nest_util.py:1006, in _tf_core_pack_sequence_as(structure, flat_sequence, expand_composites, sequence_fn)
   1003   return flat_sequence[0]
   1005 try:
-> 1006   final_index, packed = _tf_core_packed_nest_with_indices(
   1007       structure, flat_sequence, 0, is_nested_fn, sequence_fn
   1008   )
   1009   if final_index < len(flat_sequence):
   1010     raise IndexError

File ~/mambaforge/envs//lib/python3.10/site-packages/tensorflow/python/util/nest_util.py:679, in _tf_core_packed_nest_with_indices(structure, flat, index, is_nested_fn, sequence_fn)
    675 if is_nested_fn(s):
    676   new_index, child = _tf_core_packed_nest_with_indices(
    677       s, flat, index, is_nested_fn, sequence_fn
    678   )
--> 679   packed.append(sequence_fn(s, child))
    680   index = new_index
    681 else:

File ~/mambaforge/envs//lib/python3.10/site-packages/tensorflow/python/util/nest_util.py:336, in sequence_like(instance, args)
    334   assert len(args) == 1
    335   spec = instance._type_spec  # pylint: disable=protected-access
--> 336   return spec._from_components(args[0])  # pylint: disable=protected-access
    337 elif _is_type_spec(instance):
    338   # Pack a CompositeTensor's components according to a TypeSpec.
    339   assert len(args) == 1

File ~/mambaforge/envs//lib/python3.10/site-packages/tensorflow/python/framework/extension_type.py:472, in ExtensionTypeSpec._from_components(self, components)
    470 spec_tuple = tuple(self.__dict__.values())
    471 components_iter = iter(components)
--> 472 flat = [
    473     next(components_iter) if isinstance(x, type_spec.TypeSpec) else x
    474     for x in nest.flatten(spec_tuple)
    475 ]
    476 if list(components_iter):
    477   raise ValueError(
    478       'Cannot build an ExtensionType instance from components '
    479       'because more components are provided than the number expected '
    480       'by the type spec.'
    481   )

File ~/mambaforge/envs//lib/python3.10/site-packages/tensorflow/python/framework/extension_type.py:473, in <listcomp>(.0)
    470 spec_tuple = tuple(self.__dict__.values())
    471 components_iter = iter(components)
    472 flat = [
--> 473     next(components_iter) if isinstance(x, type_spec.TypeSpec) else x
    474     for x in nest.flatten(spec_tuple)
    475 ]
    476 if list(components_iter):
    477   raise ValueError(
    478       'Cannot build an ExtensionType instance from components '
    479       'because more components are provided than the number expected '
    480       'by the type spec.'
    481   )

StopIteration: 
```/lib/python3.10/site-packages/tensorflow/python/framework/extension_type.py:473, in <listcomp>(.0)
    470 spec_tuple = tuple(self.__dict__.values())
    471 components_iter = iter(components)
    472 flat = [
--> 473     next(components_iter) if isinstance(x, type_spec.TypeSpec) else x
    474     for x in nest.flatten(spec_tuple)
    475 ]
    476 if list(components_iter):
    477   raise ValueError(
    478       'Cannot build an ExtensionType instance from components '
    479       'because more components are provided than the number expected '
    480       'by the type spec.'
    481   )

StopIteration:
```
