Set the global random state of :mod:`torch`, :mod:`torch.cuda`, :mod:`numpy` and Python in the current process.

Autoencoder3d adapted from Stable Diffusion 1.x, 2.x and XL.

Autoencoder3d adapted from Stable Diffusion 1.x, 2.x and XL.

NorMuon variance reduction. Algebraically fuses the normalization steps to minimize memory ops.

Perform approximate Singular Value Decomposition (SVD) using Subsampled Randomized Hadamard Transform (SRHT). Args: - input_matrix (torc

Check if parameters set are all okay and raise error if there is any strange combination.

r""" Clean memory on the specified device, will be called from training scripts.

Clips gradient norm of an iterable of parameters. This has been adapted from Nvidia megatron. We add norm averaging to consider MoE params wh

Compute appropriate scale for the given tensor and target dtype. Args: tensor: Input tensor to be quantized dtype: Target dt

Initialize distributed training environment.

Initialize distributed training environment.

r""" Args: x(Tensor): Shape [B, L1, C] e(Tensor): Shape [B, L1, C]

Returns (mantissa_bits, total_bits) for each format. mantissa_bits excludes the implicit leading 1.

Args: sigma (tensor): noise level Returns: loss weights per sigma noise level

r""" Initialize model parameters using Xavier initialization.

A context manager that resets the global random state on exit to what it was before entering. It supports isolating the states for PyTorch, Numpy,