Monkfish: Distributed latent video model training on TPUs (and other stuff maybe)

This is the training code for a 2 stage autoregressive video model.

TODO:

• Chunked scatter/gather/init functions
• Dtype conversions at scatter/gather/init functions
• Distributed data loading
• Distributed model training
• Multi-platform file backend via PyFilesystem2
• GPU Support
• Text conditional diffusion Transformer