kd.typing.ElementSpec# [[Source]] from clu.data.dataset_iterator import ArraySpec, ElementSpec, PyTree