Make DistributedDataset a compositeTensor so that it can be returned from experimental_distribute_dataset method and be used in concrete tf function.
Make DistributedDataset a compositeTensor so that it can be returned from experimental_distribute_dataset method and be used in concrete tf function.

This is mainly relevant for PS_Strategy where dataset_fn is converted into a tf_function which requires that return types are Tensors.

PiperOrigin-RevId: 360559465

keras-team/keras,keras-team/keras