Small attention clean-up: Replace asserts with ValueError; align message style.
Small attention clean-up: Replace asserts with ValueError; align message style.

PiperOrigin-RevId: 358918151

google/trax,google/trax