r"""Returns a nested state dictionary containing the whole internal dict of the module. OptimizerModules and other common data structures

Compresses sequence based on selector. NOTE: Despite the name, this function can compress both lists and tuples, but will always return a tup

This implementation requires the stored and loaded states to be fully initialized. Because of introduced strictness it allows us to:

Test helper function that verifies preconditioner updates and preconditioning. This function takes a preconditioner_list and

Compare the performance of two optimizers on a simple neural network using the same device. Args: control_optim_factory (Callable[[P

Generates accumulated pairwise indices for a given input list. This is useful for generating interval indices for iterating through a list given

Compare the performance of the same optimizer on a simple neural network across CPU and another device. Args: optim_factory (Callabl

Compare the performance of two optimizers on models across different devices by evaluating their weights and loss. Args: control_opt

Merge and block gradients. NOTE: This function MUST be called in the step function of the optimizer after the gradient has been updat

Retrieves all non-abstract (instantiable) subclasses of a given class. This function uses a helper function to recursively find all unique s

Compiles parameter metadata necessary for FSDP Shampoo. Args: module (nn.Module): Module to compile metadata for. Returns: p

Add epsilon * I to matrix (if square) or epsilon (if vector). Args: A (Tensor): Matrix of interest. epsilon (float): Value to add

Computes matrix root inverse of square symmetric positive definite matrix. Args: A (Tensor): Square matrix of interest. root (Fra

Batches an iterable into chunks of size n. Note: This is a re-implementation of itertools.batched which is available in Python 3.12+. Co

Chunks tensor across multiple dimensions based on splits. This function recursively splits a tensor along all of its dimensions using the spe

Creates a partial function with arguments from config that match func's parameter names. This function examines the fields in the config obj

Implementation of test_update_params - called by concrete test classes. This test verifies that: - use_masked_tensors=True: U

Resolve a string to a torch.dtype.

Returns device mesh from provided device type, mesh, and mesh dim names. This function will cache previous meshes according to the input. Arg

Reshapes tensor by merging small dimensions. This function merges adjacent dimensions of a tensor when their product is below the specified t

Return the optimizer state as a dict.