Partition tensor into blocks.

Computes mat_m^p, for p a positive integer. Args: mat_m: a square matrix p: a positive integer Returns: mat_m^p

Power iteration. Compute the maximum eigenvalue of mat, for scaling. v is a random vector with values in (-1, 1) Args: mat_g: the symmetri

Merge small dimensions. If there are some small dimensions, we collapse them: e.g. [1, 2, 512, 1, 2048, 1, 3, 4] --> [1024, 2048, 12] if max_dim

Compute L^{-1/exp} for each stats matrix L.

Returns exponent to use for inverse-pth root M^{-1/p}.

Initialize the PyTorch state of for a single variable.

Merge partitions back to original shape.

Precondition the gradient. Args: grad: A gradient tensor to precondition. Returns: A preconditioned gradient.

A method to compute G^{-1/p} using a coupled Newton iteration. See for example equation 3.2 on page 9 of: A Schur-Newton Method for the Matrix p-

Compute statistics from gradients and add to the correct state entries. Args: grad: Gradient to compute statistics from.

Splits larger dimensions into smaller ones, for preconditioning. For example, if a variable has shape (4096, 512), we might split the 4096 into 4

Merge the split gradients back into a single array.

Computes the shape of the variable for preconditioning. If the variable has several small dimensions, we can reshape it so that there are fewer b

Split up the gradient according to the blocking strategy.