nabling memory-efficient attention with xFormers. JAX/Flax training is also supported for efficient training on TPUs and GPUs, but it doesn’t support gradient checkpointing, gradient accumulation or xFormers. A GPU with at least 30GB of memory or a TPU v3 is recommended for training with Flax. This guide will explore the train_text_to_image.py training script to help you become familiar with it, and how you can adapt it for your own use-case. B