Depend on LazyArray in Orbax, as the functionality is the same and we will soon rely on Orbax for checkpointing.
Depend on LazyArray in Orbax, as the functionality is the same and we will soon rely on Orbax for checkpointing.

PiperOrigin-RevId: 450702331

google-research/t5x