Utilities#
The utils module contains various utility functions and helpers for the DiffuseNNX library.
Overview#
This module provides essential utilities for:
Checkpointing: Model and optimizer state saving/loading with Orbax
EMA (Exponential Moving Average): Model weight averaging for stable training
Google Cloud Integration: GCS utilities for cloud storage and data access
Initialization: Model parameter initialization and registry management
Logging: Logging utilities and configuration for training monitoring
Sharding: Distributed training utilities with JAX mesh and FSDP support
Visualization: Training monitoring and sample generation visualization
Weights & Biases: Experiment tracking integration and logging
Advanced Tutorials#
We also provide a Tutorial for our sharding implementation (cr. Georgy)!