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