Access feature shapes from dataset directly to handle >2D features in `eval`.
Access feature shapes from dataset directly to handle >2D features in `eval`.

PiperOrigin-RevId: 425717127

google-research/t5x