diff --git a/CHANGELOG.md b/CHANGELOG.md
index 31e8e92..ae5927d 100644
--- a/CHANGELOG.md
+++ b/CHANGELOG.md
@@ -1,6 +1,16 @@
 # Changelog
 
-Release notes for the OSS mirror, generated by `scripts/oss_sync.sh` from the internal source tree.
+Release notes for the OSS mirror. Seeded by `scripts/oss_sync.sh` from the internal source tree and edited for release.
+
+## 2026-08-21
+
+- [shampoo] Add per-head block splitting for attention projections, via two new `DDPDistributedConfig` fields: `head_split_size` and `head_split_param_names`. When `head_split_size > 0`, a matching 2D parameter is split along dim 0 into `head_split_size`-row blocks instead of taking the default merge-and-block path, so a matrix preconditioner (Shampoo or spectral descent) sees one block per attention head rather than the whole fused projection. A parameter is split iff its name contains any substring in `head_split_param_names`; parameters in the same group that do not match keep the default blocking. Opt-in — `head_split_size` defaults to 0, so the default path is unchanged. Requires per-parameter names under `param_names` in the param group, and raises if they are missing or their count does not match `params`; a negative `head_split_size` raises at config construction, and at block time a non-2D parameter, or a size that does not divide the parameter's first dimension, raises. If `head_split_size` is set but no name matches, a warning is logged and the split is a no-op. Implemented for the DDP distributor.
+
+## 2026-08-19
+
+- [shampoo] Add `NewtonSchulzRootInvConfig`: a matmul-only coupled Newton-Schulz iteration, usable as the `amortized_computation_config` of `RootInvShampooPreconditionerConfig` in place of the default eigendecomposition. Opt-in — `DefaultShampooConfig` is unchanged. The iteration runs a fixed number of steps with no residual-based stopping criterion, so it introduces no host-device synchronization. Only power-of-two inverse roots are supported (blocks of order 1, 2, or 4); anything else raises at optimizer construction rather than degrading to a per-factor-matrix warning that reuses a stale preconditioner. `relative_epsilon` (default `1e-6`) floors the ridge at a fraction of `|A|_F` and is applied unconditionally, so on the rank-deficient factor matrices seen early in training this path regularizes more aggressively than eigendecomposition at the same epsilon and the two do not agree there. `coefficients` takes a per-iteration schedule of `(a, b, c)` triples for `p(x) = a x + b x^3 + c x^5`, defaulting to a 10-step Polar Express schedule; tf32 is disabled inside the iteration by default (`disable_tf32=True`).
+- [shampoo] Annotate `ShampooPT2CompileConfig`'s `make_dataclass`-synthesized base class `Any` so mypy accepts it as a base class, and add test coverage asserting the synthesized field set matches `torch.compile`'s signature (no runtime change).
+- [shampoo] Remove unused type-error suppression comments across the optimizer, examples, and tests (no behavior change).
 
 ## 2026-07-24
 
diff --git a/distributed_shampoo/__init__.py b/distributed_shampoo/__init__.py
index 767cf6a..2ec751d 100644
--- a/distributed_shampoo/__init__.py
+++ b/distributed_shampoo/__init__.py
@@ -23,6 +23,7 @@
     EighEigendecompositionConfig,
     MatrixFunctionConfig,
     NewtonSchulzOrthogonalizationConfig,
+    NewtonSchulzRootInvConfig,
     OrthogonalizationConfig,
     PerturbationConfig,
     PseudoInverseConfig,
@@ -118,6 +119,7 @@
     "DefaultEigenConfig",  # Default `RootInvConfig` using `EigenConfig`.
     "CoupledNewtonConfig",  # Based on `RootInvConfig`.
     "CoupledHigherOrderConfig",  # Based on `RootInvConfig`.
+    "NewtonSchulzRootInvConfig",  # Based on `RootInvConfig`.
     "OrthogonalizationConfig",  # Abstract base class (based on `MatrixFunctionConfig`).
     "SVDOrthogonalizationConfig",  # Based on `OrthogonalizationConfig`.
     "NewtonSchulzOrthogonalizationConfig",  # Based on `OrthogonalizationConfig`.
diff --git a/distributed_shampoo/distributed_shampoo.py b/distributed_shampoo/distributed_shampoo.py
index 0bbbc84..bdf8014 100644
--- a/distributed_shampoo/distributed_shampoo.py
+++ b/distributed_shampoo/distributed_shampoo.py
@@ -981,7 +981,6 @@ def _instantiate_per_group_step(
         # Use PT2 to compile the step function for each parameter group.
         self._per_group_step: Callable[..., None] = (
             torch.compile(
-                # pyrefly: ignore [bad-argument-type]
                 self._per_group_step_impl,
                 # pyrefly: ignore [bad-argument-type]
                 **asdict(shampoo_pt2_compile_config),
@@ -1370,7 +1369,7 @@ def _apply_in_place_primal_averaging(
         # This computes: - (1 - mu_x * mu_y) * lr * P.
         torch._foreach_mul_(
             masked_blocked_search_directions,
-            (1 - train_interp_coeff * eval_interp_coeff),  # type: ignore
+            (1 - train_interp_coeff * eval_interp_coeff),
         )
         # This computes: (1 - mu_x) * (Z_old - Y) - (1 - mu_x * mu_y) * lr * P.
         # pyrefly: ignore [no-matching-overload]
diff --git a/distributed_shampoo/distributor/shampoo_ddp_distributor.py b/distributed_shampoo/distributor/shampoo_ddp_distributor.py
index df38642..5acab4a 100644
--- a/distributed_shampoo/distributor/shampoo_ddp_distributor.py
+++ b/distributed_shampoo/distributor/shampoo_ddp_distributor.py
@@ -24,6 +24,7 @@
 from distributed_shampoo.shampoo_types import (
     DDPDistributedConfig,
     DISTRIBUTED_CONFIG,
+    PARAM_NAMES,
     PARAMS,
     ShampooRuntimeConfig,
 )
@@ -220,6 +221,47 @@ def __init__(
             group_rank=group_rank,
         )
 
+    def _resolve_head_split_sizes(self) -> tuple[int | None, ...] | None:
+        """Resolve per-parameter head-split sizes from the DDP config.
+
+        Returns None (head-split disabled) unless ``head_split_size > 0``. When
+        enabled, a parameter is head-split iff its name contains any substring in
+        ``head_split_param_names``; ``head_split_size`` (head_dim) is used for
+        those and None for the rest. Requires per-parameter names under
+        PARAM_NAMES in the param group. Works with any matrix preconditioner
+        (e.g. Shampoo or spectral descent / Muon).
+        """
+        distributed_config: DDPDistributedConfig = self._param_group[DISTRIBUTED_CONFIG]
+        head_split_size = distributed_config.head_split_size
+        if not head_split_size:
+            return None
+        if PARAM_NAMES not in self._param_group:
+            raise ValueError(
+                f"head_split_size requires per-parameter names under "
+                f"'{PARAM_NAMES}' in the param group."
+            )
+        param_names = self._param_group[PARAM_NAMES]
+        num_params = len(self._param_group[PARAMS])
+        if len(param_names) != num_params:
+            raise ValueError(
+                f"len({PARAM_NAMES})={len(param_names)} must equal "
+                f"len(params)={num_params}."
+            )
+        head_split_sizes = tuple(
+            head_split_size
+            if any(
+                pattern in name for pattern in distributed_config.head_split_param_names
+            )
+            else None
+            for name in param_names
+        )
+        if not any(size is not None for size in head_split_sizes):
+            logger.warning(
+                "head_split_size is set but no parameter name matched "
+                "head_split_param_names; head-split is a no-op."
+            )
+        return head_split_sizes
+
     @overload
     @torch.no_grad()
     def _get_params_or_grads(
diff --git a/distributed_shampoo/distributor/shampoo_distributor.py b/distributed_shampoo/distributor/shampoo_distributor.py
index 645901d..4bcfd8b 100644
--- a/distributed_shampoo/distributor/shampoo_distributor.py
+++ b/distributed_shampoo/distributor/shampoo_distributor.py
@@ -7,8 +7,9 @@
 
 """
 
+import logging
 from abc import ABC, abstractmethod
-from collections.abc import Iterable
+from collections.abc import Callable, Iterable
 from functools import partial
 from operator import attrgetter
 from typing import Any, Literal, overload
@@ -29,6 +30,8 @@
 )
 from torch import Tensor
 
+logger: logging.Logger = logging.getLogger(__name__)
+
 
 ###### DISTRIBUTOR CLASSES ######
 class DistributorInterface(ABC):
@@ -51,6 +54,13 @@ def __init__(
         self._runtime_config: ShampooRuntimeConfig = (
             runtime_config if runtime_config is not None else ShampooRuntimeConfig()
         )
+        # Per-parameter head-split sizes (None disables head-split for every
+        # param). Resolved via virtual dispatch so only distributors that
+        # support head-split (currently DDP) return a non-None spec. Must be set
+        # before _merge_and_block_parameters() so blocking can consult it.
+        self._head_split_sizes: tuple[int | None, ...] | None = (
+            self._resolve_head_split_sizes()
+        )
         # Merge and block parameters creates self._global_blocked_params and self._global_num_blocks_per_param
         # Global blocked params are all the blocked parameters after merging and blocking.
         # Global num blocks per param stores the number of blocks for each global parameter.
@@ -75,6 +85,21 @@ def __init__(
         # Local block info list contains information about each block masked by the distributor selector.
         self._local_block_info_list: tuple[BlockInfo, ...]
 
+    def _resolve_head_split_sizes(self) -> tuple[int | None, ...] | None:
+        """Return per-parameter head-split sizes, or None to disable head-split.
+
+        Base implementation disables head-split. Distributors that support it
+        (currently only ``DDPDistributor``) override this. A warning is emitted
+        if head-split was requested on an unsupported distributor so the config
+        is not silently ignored.
+        """
+        if getattr(self._param_group[DISTRIBUTED_CONFIG], "head_split_size", 0):
+            logger.warning(
+                "head_split_size is set but ignored: head-split is only "
+                f"supported by DDPDistributor, not {type(self).__name__}."
+            )
+        return None
+
     @abstractmethod
     @torch.no_grad()
     def update_params(
@@ -170,6 +195,43 @@ def _get_params_or_grads(self, get_grad: bool = False) -> Iterable[Tensor | None
             else self._param_group[PARAMS]
         )
 
+    def _blocks_within_tensor(
+        self,
+        tensor: Tensor,
+        param_index: int,
+        merge_dims: Callable[..., tuple[int, ...]],
+    ) -> tuple[Tensor, ...]:
+        """Split one param/grad tensor into preconditioner blocks.
+
+        Default path: merge small dims, then block by ``max_preconditioner_dim``.
+        When head-split is enabled for this param, instead split dim 0 into
+        ``head_dim``-sized 2D blocks (one per attention head) and skip merging /
+        blocking so the model (input) dimension is preserved. Both the parameter
+        and gradient paths call this so their block decomposition stays
+        identical.
+        """
+        head_dim = (
+            None
+            if self._head_split_sizes is None
+            else self._head_split_sizes[param_index]
+        )
+        if head_dim is not None:
+            if tensor.dim() != 2:
+                raise ValueError(
+                    "head-split requires a 2D parameter, got shape "
+                    f"{tuple(tensor.shape)} at param index {param_index}."
+                )
+            if tensor.shape[0] % head_dim != 0:
+                raise ValueError(
+                    f"head_split_size={head_dim} must divide out_features="
+                    f"{tensor.shape[0]} at param index {param_index}."
+                )
+            return torch.split(tensor, head_dim, dim=0)
+        return multi_dim_split(
+            tensor.view(merge_dims(tensor_shape=tensor.size())),
+            self._param_group[MAX_PRECONDITIONER_DIM],
+        )
+
     @torch.no_grad()
     def _merge_and_block_with_params(
         self, params: Iterable[Tensor]
@@ -199,11 +261,10 @@ def _merge_and_block_with_params(
             ].target_parameter_dimensionality,
         )
 
-        for param in params:
-            # Obtain blocks for each parameter after merging.
-            blocks_within_param = multi_dim_split(
-                param.view(merge_dims(tensor_shape=param.size())),
-                self._param_group[MAX_PRECONDITIONER_DIM],
+        for param_index, param in enumerate(params):
+            # Obtain blocks for each parameter after merging (or per-head split).
+            blocks_within_param = self._blocks_within_tensor(
+                param, param_index, merge_dims
             )
 
             # Generate and extend blocked parameters list.
@@ -256,11 +317,17 @@ def _merge_and_block_gradients(
             ].target_parameter_dimensionality,
         )
 
-        for grad, num_blocks, (block_index, next_block_index) in zip(
-            self._get_params_or_grads(get_grad=True),
-            self._global_num_blocks_per_param,
-            generate_pairwise_indices(self._global_num_blocks_per_param),
-            strict=True,
+        for param_index, (
+            grad,
+            num_blocks,
+            (block_index, next_block_index),
+        ) in enumerate(
+            zip(
+                self._get_params_or_grads(get_grad=True),
+                self._global_num_blocks_per_param,
+                generate_pairwise_indices(self._global_num_blocks_per_param),
+                strict=True,
+            )
         ):
             param_distributor_selector = self._distributor_selector[
                 block_index:next_block_index
@@ -283,10 +350,9 @@ def _merge_and_block_gradients(
                     f"Encountered gradient containing NaN/Inf in parameter with shape {attrgetter('shape')(grad)}. Check your model for numerical instability or consider gradient clipping."
                 )
 
-            # Obtain blocks for each gradient after merging.
-            blocks_within_grad = multi_dim_split(
-                grad.view(merge_dims(tensor_shape=grad.size())),
-                self._param_group[MAX_PRECONDITIONER_DIM],
+            # Obtain blocks for each gradient after merging (or per-head split).
+            blocks_within_grad = self._blocks_within_tensor(
+                grad, param_index, merge_dims
             )
             # Generate block-to-parameter metadata and extend blocked parameters list.
             local_masked_blocked_grads.extend(
diff --git a/distributed_shampoo/examples/utils.py b/distributed_shampoo/examples/utils.py
index a237c93..291a294 100644
--- a/distributed_shampoo/examples/utils.py
+++ b/distributed_shampoo/examples/utils.py
@@ -73,7 +73,6 @@ def get_distributed_env() -> tuple[int, int, int]:
 
 def set_seed(seed: int) -> None:
     torch.manual_seed(seed)
-    # pyrefly: ignore [bad-argument-type]
     np.random.seed(seed)
     random.seed(seed)
     torch.use_deterministic_algorithms(True)
diff --git a/distributed_shampoo/preconditioner/matrix_functions.py b/distributed_shampoo/preconditioner/matrix_functions.py
index 48347e4..a67c8de 100644
--- a/distributed_shampoo/preconditioner/matrix_functions.py
+++ b/distributed_shampoo/preconditioner/matrix_functions.py
@@ -31,6 +31,7 @@
     EigendecompositionConfig,
     EighEigendecompositionConfig,
     NewtonSchulzOrthogonalizationConfig,
+    NewtonSchulzRootInvConfig,
     OrthogonalizationConfig,
     PerturbationConfig,
     PseudoInverseConfig,
@@ -436,6 +437,128 @@ def matrix_inverse_root_newton(
 
         return X, M, termination_flag, iteration, error
 
+    def matrix_inverse_root_newton_schulz(
+        A: Tensor,
+        root: int,
+        coefficients: list[list[float]],
+        epsilon: float = 0.0,
+        relative_epsilon: float = 1e-6,
+        disable_tf32: bool = True,
+    ) -> Tensor:
+        """Compute matrix inverse root using the coupled Newton-Schulz iteration.
+
+        A single pass implements SqrtInverseNewtonSchulz, which returns both A^{1/2} and A^{-1/2}:
+
+            alpha <- |A|_F
+            Y <- A / alpha, Z <- I
+            for each (a, b, c) in the coefficient schedule
+                T <- Z Y
+                B <- b T + c T^2
+                Y <- a Y + Y B
+                Z <- a Z + B Z
+            A^{1/2} ~= sqrt(alpha) Y, A^{-1/2} ~= Z / sqrt(alpha)
+
+        Y and Z are updated as a coupled pair rather than by recomputing a residual from A, which is
+        what makes the iteration numerically stable; the uncoupled form diverges within ~15 iterations
+        even in double precision.
+
+        Only roots that are powers of two are supported. Since a pass yields both the square root and
+        the inverse square root, A^{-1/2^m} is obtained by chaining m passes: each of the first m - 1
+        passes feeds its square root output forward, and the last pass returns its inverse square root
+        output.
+
+        NOTE: Coefficients summing to 1 give the scalar map a fixed point at 1, so the iteration
+            converges to the inverse square root. Coefficients that do not sum to 1 -- such as the Muon
+            coefficients (3.4445, -4.7750, 2.0315) used by newton_schulz for orthogonalization, where
+            only the sign of the singular values matters -- instead converge to a band around it.
+            NewtonSchulzRootInvConfig warns about this at construction time.
+
+        NOTE: Unlike the eigendecomposition path, this iteration cannot stabilize a rank-deficient or
+            indefinite input, because it never forms the spectrum it would need to shift. Shampoo
+            factor matrices are rank-deficient early in training and pick up small negative eigenvalues
+            from floating point error, on which Z diverges. relative_epsilon guards against this by
+            flooring the ridge at a fraction of |A|_F, which dominates those spurious eigenvalues and
+            bounds the condition number. The consequence is that on rank-deficient input this function
+            regularizes more aggressively than the eigendecomposition path does for the same epsilon,
+            and therefore does not agree with it there.
+
+        References:
+            - https://arxiv.org/abs/2505.16932 (Polar Express)
+            - https://docs.modula.systems/algorithms/newton-schulz/
+
+        Args:
+            A (Tensor): Matrix of interest.
+            root (int): Root of interest. Must be a power of two.
+            coefficients (list[list[float]]): Per-iteration schedule of (a, b, c) coefficient
+                triples for the odd polynomial p(x) = a x + b x^3 + c x^5 driving the iteration. The
+                number of iterations is len(coefficients). Validated by
+                NewtonSchulzRootInvConfig.__post_init__, which is the only supported way to reach
+                this function.
+            epsilon (float): Adds epsilon * I to matrix before taking matrix root. (Default: 0.0)
+            relative_epsilon (float): Adds relative_epsilon * |A|_F * I to the matrix before taking the
+                matrix root, taking the larger of this and epsilon. Required for rank-deficient input;
+                see the note above. (Default: 1e-6)
+            disable_tf32 (bool): Whether to disable tf32 matmuls or not internally. Highly recommend
+                keeping True; tf32 cannot represent the relative_epsilon-scale eigenvalues this
+                iteration must resolve, and Z @ Y is not a Gram product, so the resulting error is
+                unstructured and diverges. (Default: True)
+
+        Returns:
+            X (Tensor): Inverse root of matrix A.
+
+        Raises:
+            ValueError: If root is not a power of two.
+
+        """
+        # This should not be reachable: RootInvShampooPreconditionerList already rejects roots that
+        # are not powers of two at optimizer construction time. Kept as a defensive guard in case
+        # this function is ever called directly, bypassing that validation.
+        if root < 2 or root & (root - 1):
+            raise ValueError(
+                f"{root=} must be a power of two to use Newton-Schulz iteration!"
+            )
+
+        tf32_flag = torch.backends.cuda.matmul.allow_tf32
+        if disable_tf32:
+            torch.backends.cuda.matmul.allow_tf32 = False
+
+        try:
+            # Add regularization, floored relative to |A|_F so that the input is positive definite. The
+            # ridge is kept as a 0-d tensor rather than a float so that no host-device synchronization
+            # is introduced.
+            identity = torch.eye(A.shape[0], dtype=A.dtype, device=A.device)
+            A = torch.addcmul(
+                A,
+                identity,
+                torch.linalg.matrix_norm(A).mul_(relative_epsilon).clamp_min_(epsilon),
+            )
+
+            # A^{-1/root} = ((...(A^{1/2})^{1/2}...)^{1/2})^{-1/2} with log2(root) nested square roots.
+            for pass_index in range(root.bit_length() - 1, 0, -1):
+                # Normalize so that the spectrum of Y lies in (0, 1]; |A|_2 <= |A|_F.
+                alpha = torch.linalg.matrix_norm(A).clamp_min(1e-8)
+                Y = A / alpha
+                # Cloned because the final pass returns Z.div_(sqrt_alpha), which mutates in place,
+                # and identity is reused across passes. torch.addmm below is out-of-place and would
+                # rebind Z on its own, but that only holds while the schedule is non-empty.
+                Z = identity.clone()
+
+                for a, b, c in coefficients:
+                    T = Z @ Y
+                    B = torch.addmm(T, T, T, beta=b, alpha=c)
+                    Y = torch.addmm(Y, Y, B, beta=a, alpha=1)
+                    Z = torch.addmm(Z, B, Z, beta=a, alpha=1)
+
+                sqrt_alpha = alpha.sqrt()
+                # Feed A^{1/2} into the next pass; the final pass produces the inverse root.
+                A = Y.mul_(sqrt_alpha) if pass_index > 1 else Z.div_(sqrt_alpha)
+        finally:
+            # Always restore tf32 mode unconditionally, so we skip the disable_tf32 check. When
+            # disable_tf32=False, this is a no-op since tf32_flag already equals the current value.
+            torch.backends.cuda.matmul.allow_tf32 = tf32_flag
+
+        return A
+
     def matrix_inverse_root_higher_order(
         A: Tensor,
         root: Fraction,
@@ -679,6 +802,16 @@ def matrix_inverse_root_higher_order(
                 logger.warning(
                     "Newton did not converge and reached maximum number of iterations!"
                 )
+        case NewtonSchulzRootInvConfig():
+            # NOTE: Use Fraction.is_integer() instead when downstream applications are Python 3.12+ available
+            if root.denominator != 1:
+                raise ValueError(
+                    f"{root.denominator=} must be equal to 1 to use Newton-Schulz iteration!"
+                )
+
+            X = _assign_function_args_from_config(
+                func=matrix_inverse_root_newton_schulz, config=root_inv_config
+            )(A=A, root=root.numerator, epsilon=epsilon)
         case CoupledHigherOrderConfig():
             X, _, termination_flag, _, _ = _assign_function_args_from_config(
                 func=matrix_inverse_root_higher_order, config=root_inv_config
diff --git a/distributed_shampoo/preconditioner/matrix_functions_types.py b/distributed_shampoo/preconditioner/matrix_functions_types.py
index 02e7c63..ade8cf3 100644
--- a/distributed_shampoo/preconditioner/matrix_functions_types.py
+++ b/distributed_shampoo/preconditioner/matrix_functions_types.py
@@ -7,11 +7,15 @@
 
 """
 
+import logging
+import math
 from collections.abc import Callable
 from dataclasses import dataclass, field
 
 from distributed_shampoo.utils.abstract_dataclass import AbstractDataclass
 
+logger: logging.Logger = logging.getLogger(__name__)
+
 
 @dataclass(init=False)
 class RankDeficientStabilityConfig(AbstractDataclass):
@@ -213,6 +217,86 @@ class CoupledNewtonConfig(RootInvConfig):
     tolerance: float = 1e-6
 
 
+@dataclass(kw_only=True)
+class NewtonSchulzRootInvConfig(RootInvConfig):
+    """Configuration for matrix root inverse via the coupled Newton-Schulz iteration.
+
+    Unlike CoupledNewtonConfig and CoupledHigherOrderConfig, the iteration runs for a fixed number of
+    steps and never evaluates a residual-based stopping criterion, so it incurs no host-device
+    synchronization and is expressed entirely as matmuls. Only roots that are powers of two are
+    supported, i.e. blocks of order 1, 2, and 4.
+
+    WARNING: On rank-deficient input this regularizes more aggressively than the eigendecomposition
+        path does for the same epsilon, so the two do not agree there. See relative_epsilon.
+
+    Attributes:
+        relative_epsilon (float): Floors the ridge added before the iteration at
+            relative_epsilon * |A|_F. Unlike the eigendecomposition path, this iteration cannot
+            stabilize a rank-deficient or indefinite matrix, and Shampoo factor matrices are
+            rank-deficient early in training; without this floor the iteration diverges to NaN.
+            (Default: 1e-6)
+        coefficients (list[list[float]]): Per-iteration schedule of (a, b, c) coefficient triples
+            for the odd polynomial p(x) = a x + b x^3 + c x^5 driving the iteration. The number
+            of iterations is len(coefficients).
+            (Default: Polar Express 10-step schedule from ASGO.)
+        disable_tf32 (bool): Whether to disable tf32 matmuls or not internally. Highly recommend
+            keeping True. The iteration is built entirely out of matmuls and must resolve
+            eigenvalues down to relative_epsilon, which tf32's 10-bit mantissa cannot represent;
+            unlike the factor matrix accumulation, Z @ Y is not a Gram product, so tf32 error there
+            is unstructured and drives the iteration to NaN. (Default: True)
+
+    """
+
+    @staticmethod
+    def _get_default_coefficients() -> list[list[float]]:
+        return [
+            [8.28721201814563, -23.595886519098837, 17.300387312530933],
+            [4.107059111542203, -2.9478499167379106, 0.5448431082926601],
+            [3.9486908534822946, -2.9089021159629490, 0.5518191394370137],
+            [3.3184196573706015, -2.4884880243148740, 0.5100489401237200],
+            [2.300652019954817, -1.6689039845747493, 0.4188073119525673],
+            [1.891301407787398, -1.2679958271945868, 0.3768040894852483],
+            [1.8750014808534479, -1.2500016453999487, 0.3750001645474248],
+            [1.875, -1.25, 0.375],
+            [1.875, -1.25, 0.375],
+            [1.875, -1.25, 0.375],
+        ]
+
+    relative_epsilon: float = 1e-6
+    disable_tf32: bool = True
+    # TODO: Clean up coefficient definition -- consider using list[tuple[float, float, float]]
+    # to enforce 3-tuples, and define defaults from the training pipeline side.
+    coefficients: list[list[float]] = field(default_factory=_get_default_coefficients)
+
+    def __post_init__(self) -> None:
+        if len(self.coefficients) == 0:
+            raise ValueError("coefficients must be non-empty.")
+        for index, entry in enumerate(self.coefficients):
+            if len(entry) != 3:
+                raise ValueError(
+                    f"coefficients[{index}] must contain exactly three coefficients (a, b, c) for "
+                    f"p(x) = a x + b x^3 + c x^5, but {entry=} has {len(entry)}."
+                )
+            for coefficient in entry:
+                if isinstance(coefficient, bool) or not isinstance(
+                    coefficient, (int, float)
+                ):
+                    raise ValueError(
+                        f"coefficients[{index}] must contain real numbers, but {entry=} contains "
+                        f"{coefficient!r} of type {type(coefficient).__name__}."
+                    )
+                if not math.isfinite(coefficient):
+                    raise ValueError(
+                        f"coefficients[{index}] must be finite, but {entry=} contains {coefficient}."
+                    )
+        final_coefficients = self.coefficients[-1]
+        if not math.isclose(sum(final_coefficients), 1.0):
+            logger.warning(
+                f"{final_coefficients=} do not sum to 1, so the Newton-Schulz iteration has no fixed "
+                "point at 1 and will converge to a band around the inverse root rather than to it."
+            )
+
+
 @dataclass(kw_only=True)
 class CoupledHigherOrderConfig(RootInvConfig):
     """Configuration for matrix root inverse via coupled higher-order method.
diff --git a/distributed_shampoo/preconditioner/shampoo_preconditioner_list.py b/distributed_shampoo/preconditioner/shampoo_preconditioner_list.py
index d4a6732..db57717 100644
--- a/distributed_shampoo/preconditioner/shampoo_preconditioner_list.py
+++ b/distributed_shampoo/preconditioner/shampoo_preconditioner_list.py
@@ -29,6 +29,7 @@
 from distributed_shampoo.preconditioner.matrix_functions_types import (
     EigendecompositionConfig,
     MatrixFunctionConfig,
+    NewtonSchulzRootInvConfig,
     RootInvConfig,
 )
 from distributed_shampoo.preconditioner.preconditioner_list import (
@@ -686,6 +687,26 @@ def __post_init__(self) -> None:
             == len(self.factor_matrices)
             == len(self.inv_factor_matrices)
         )
+        # Fail fast rather than at the first amortized computation: an input the amortized
+        # computation cannot handle would otherwise surface as a swallowed per-factor-matrix warning
+        # that silently reuses the stale preconditioner until
+        # num_tolerated_failed_amortized_computations is hit.
+        if isinstance(self.amortized_computation_config, NewtonSchulzRootInvConfig):
+            if unsupported_roots := sorted(
+                {
+                    root
+                    for root in self.roots
+                    if not float(root).is_integer()
+                    or int(root) < 2
+                    or int(root) & (int(root) - 1)
+                }
+            ):
+                raise ValueError(
+                    f"{type(self.amortized_computation_config).__name__} only supports inverse roots "
+                    f"that are powers of two, but {unsupported_roots=} were requested. Merge or block "
+                    "the offending parameters down to order 1, 2, or 4, or set inverse_exponent_override "
+                    "to the reciprocal of a power of two."
+                )
 
 
 @dataclass(kw_only=True)
diff --git a/distributed_shampoo/preconditioner/tests/matrix_functions_test.py b/distributed_shampoo/preconditioner/tests/matrix_functions_test.py
index 0ccf5b7..296e051 100644
--- a/distributed_shampoo/preconditioner/tests/matrix_functions_test.py
+++ b/distributed_shampoo/preconditioner/tests/matrix_functions_test.py
@@ -8,6 +8,7 @@
 """
 
 import itertools
+import math
 import re
 import unittest
 from collections.abc import Callable
@@ -37,6 +38,7 @@
     EigendecompositionConfig,
     EighEigendecompositionConfig,
     NewtonSchulzOrthogonalizationConfig,
+    NewtonSchulzRootInvConfig,
     OrthogonalizationConfig,
     PerturbationConfig,
     PseudoInverseConfig,
@@ -661,6 +663,207 @@ def A_tridiagonal_2(n: int, alpha: float, beta: float) -> Tensor:
         )
 
 
+@instantiate_parametrized_tests
+class NewtonSchulzRootInverseTest(unittest.TestCase):
+    @staticmethod
+    def _spd_matrix(n: int, condition_number: float) -> Tensor:
+        torch.manual_seed(42)
+        Q, _ = torch.linalg.qr(torch.randn(n, n, dtype=torch.float64))
+        eigenvalues = torch.logspace(
+            0, -math.log10(condition_number), n, dtype=torch.float64
+        )
+        return ((Q * eigenvalues) @ Q.T).float()
+
+    @staticmethod
+    def _relative_error(X: Tensor, expected: Tensor) -> float:
+        """Normwise relative error. Elementwise rtol is not meaningful here: entries of the inverse
+        root that are near zero carry large relative error at negligible absolute error."""
+        return (
+            torch.dist(X, expected, p=torch.inf)
+            / torch.linalg.norm(expected, ord=torch.inf)
+        ).item()
+
+    @parametrize("root", [2, 4, 8])
+    @parametrize("n", [10, 100])
+    def test_newton_schulz_root_inverse_identity(self, n: int, root: int) -> None:
+        torch.testing.assert_close(
+            matrix_inverse_root(
+                A=torch.eye(n),
+                root=Fraction(root),
+                root_inv_config=NewtonSchulzRootInvConfig(),
+            ),
+            torch.eye(n),
+            atol=1e-5,
+            rtol=1e-5,
+        )
+
+    @parametrize("root", [2, 4, 8])
+    # Attainable accuracy is floored by the condition number, not by the iteration count.
+    @parametrize("condition_number, tolerance", [(1e2, 1e-5), (1e4, 1e-4), (1e6, 1e-2)])
+    def test_newton_schulz_root_inverse_matches_eigen(
+        self, condition_number: float, tolerance: float, root: int
+    ) -> None:
+        A = NewtonSchulzRootInverseTest._spd_matrix(
+            n=64, condition_number=condition_number
+        )
+        self.assertLessEqual(
+            NewtonSchulzRootInverseTest._relative_error(
+                # relative_epsilon is disabled so this measures the accuracy of the iteration itself
+                # rather than the extra regularization it applies by default.
+                matrix_inverse_root(
+                    A=A,
+                    root=Fraction(root),
+                    root_inv_config=NewtonSchulzRootInvConfig(relative_epsilon=0.0),
+                ),
+                matrix_inverse_root(
+                    A=A, root=Fraction(root), root_inv_config=EigenConfig()
+                ),
+            ),
+            tolerance,
+        )
+
+    def test_newton_schulz_root_inverse_coefficient_schedule(self) -> None:
+        """A per-iteration coefficient schedule, the form Polar Express supplies, is accepted."""
+        A = NewtonSchulzRootInverseTest._spd_matrix(n=32, condition_number=1e4)
+        X = matrix_inverse_root(
+            A=A,
+            root=Fraction(4),
+            root_inv_config=NewtonSchulzRootInvConfig(
+                coefficients=[[3.0, -16.0 / 5.0, 6.0 / 5.0]] * 12
+            ),
+        )
+        self.assertTrue(torch.isfinite(X).all())
+
+    def test_newton_schulz_root_inverse_applies_epsilon(self) -> None:
+        # Singular matrix: without the epsilon ridge the inverse root does not exist.
+        A = torch.tensor([[1.0, 0.0], [0.0, 0.0]])
+        epsilon = 1e-2
+        torch.testing.assert_close(
+            matrix_inverse_root(
+                A=A,
+                root=Fraction(2),
+                root_inv_config=NewtonSchulzRootInvConfig(),
+                epsilon=epsilon,
+            ),
+            matrix_inverse_root(
+                A=A,
+                root=Fraction(2),
+                root_inv_config=EigenConfig(),
+                epsilon=epsilon,
+            ),
+            atol=1e-3,
+            rtol=1e-3,
+        )
+
+    @parametrize("root", [1, 3, 6])
+    def test_newton_schulz_root_inverse_unsupported_root(self, root: int) -> None:
+        self.assertRaisesRegex(
+            ValueError,
+            re.escape(
+                f"root={root} must be a power of two to use Newton-Schulz iteration!"
+            ),
+            matrix_inverse_root,
+            A=torch.eye(2),
+            root=Fraction(root),
+            root_inv_config=NewtonSchulzRootInvConfig(),
+        )
+
+    def test_newton_schulz_root_inverse_non_integer_root(self) -> None:
+        self.assertRaisesRegex(
+            ValueError,
+            re.escape(
+                "root.denominator=3 must be equal to 1 to use Newton-Schulz iteration!"
+            ),
+            matrix_inverse_root,
+            A=torch.tensor([[1.0, 0.0], [0.0, 4.0]]),
+            root=Fraction(2, 3),
+            root_inv_config=NewtonSchulzRootInvConfig(),
+        )
+
+    def test_newton_schulz_root_inverse_empty_coefficients(self) -> None:
+        self.assertRaisesRegex(
+            ValueError,
+            re.escape("coefficients must be non-empty."),
+            NewtonSchulzRootInvConfig,
+            coefficients=[],
+        )
+
+    def test_newton_schulz_root_inverse_warns_on_coefficients_not_summing_to_one(
+        self,
+    ) -> None:
+        with self.assertLogs(level="WARNING") as cm:
+            NewtonSchulzRootInvConfig(coefficients=[[3.4445, -4.7750, 2.0315]])
+        self.assertIn("do not sum to 1", "".join(r.msg for r in cm.records))
+
+    @parametrize(
+        "coefficients",
+        [
+            [[1.875, -1.25]],
+            [[1.875, -1.25, 0.375, 0.0]],
+            [[1.875, -1.25, 0.375], [1.0, 0.0]],
+        ],
+    )
+    def test_newton_schulz_root_inverse_coefficients_wrong_arity(
+        self, coefficients: list[list[float]]
+    ) -> None:
+        self.assertRaisesRegex(
+            ValueError,
+            re.escape("must contain exactly three coefficients"),
+            NewtonSchulzRootInvConfig,
+            coefficients=coefficients,
+        )
+
+    @parametrize("coefficient", [float("nan"), float("inf"), float("-inf")])
+    def test_newton_schulz_root_inverse_non_finite_coefficients(
+        self, coefficient: float
+    ) -> None:
+        self.assertRaisesRegex(
+            ValueError,
+            re.escape("must be finite"),
+            NewtonSchulzRootInvConfig,
+            coefficients=[[1.875, -1.25, coefficient]],
+        )
+
+    @parametrize("coefficient", ["0.375", None, True])
+    def test_newton_schulz_root_inverse_non_numeric_coefficients(
+        self, coefficient: object
+    ) -> None:
+        self.assertRaisesRegex(
+            ValueError,
+            re.escape("must contain real numbers"),
+            NewtonSchulzRootInvConfig,
+            coefficients=[[1.875, -1.25, coefficient]],
+        )
+
+    def test_newton_schulz_root_inverse_accepts_integer_coefficients(self) -> None:
+        """Integers are real numbers; the arity/finiteness checks must not reject them."""
+        self.assertTrue(
+            torch.isfinite(
+                matrix_inverse_root(
+                    A=torch.eye(4),
+                    root=Fraction(2),
+                    root_inv_config=NewtonSchulzRootInvConfig(
+                        coefficients=[[3, -3, 1]]
+                    ),
+                )
+            ).all()
+        )
+
+    @parametrize("dtype", [torch.float32, torch.float64])
+    def test_newton_schulz_root_inverse_accepts_supported_dtype(
+        self, dtype: torch.dtype
+    ) -> None:
+        X = matrix_inverse_root(
+            A=NewtonSchulzRootInverseTest._spd_matrix(n=16, condition_number=1e2).to(
+                dtype=dtype
+            ),
+            root=Fraction(2),
+            root_inv_config=NewtonSchulzRootInvConfig(),
+        )
+        self.assertIs(X.dtype, dtype)
+        self.assertTrue(torch.isfinite(X).all())
+
+
 class CoupledHigherOrderRootInverseTest(unittest.TestCase):
     def test_root_with_big_numerator_denominator(self) -> None:
         A = torch.tensor([[1.0, 0.0], [0.0, 4.0]])
diff --git a/distributed_shampoo/shampoo_types.py b/distributed_shampoo/shampoo_types.py
index 69f8864..caaa37d 100644
--- a/distributed_shampoo/shampoo_types.py
+++ b/distributed_shampoo/shampoo_types.py
@@ -12,6 +12,7 @@
 from collections.abc import Callable
 from dataclasses import dataclass, field, make_dataclass
 from inspect import signature
+from typing import Any
 
 import torch
 from distributed_shampoo.preconditioner.matrix_functions_types import (
@@ -50,6 +51,7 @@
 LR = "lr"
 MAX_PRECONDITIONER_DIM = "max_preconditioner_dim"
 PARAMS = "params"  # While this is stored in groups by default, we do not checkpoint this quantity.
+PARAM_NAMES = "param_names"
 PEAK_LR = "peak_lr"
 PRECONDITION_FREQUENCY = "precondition_frequency"
 PRECONDITIONER_CONFIG = "preconditioner_config"
@@ -1085,6 +1087,12 @@ class DDPDistributedConfig(DistributedConfig):
     communication_dtype: torch.dtype = torch.float32
     num_trainers_per_group: int = -1
     communicate_params: bool = False
+    # Per-head splitting of attention QKV weights for Muon / spectral descent.
+    # head_split_size = head_dim (0 disables). head_split_param_names lists the
+    # FQN substrings (e.g. "qkv_linear.") of the projections to split along
+    # dim 0 into one 2D [head_dim, in_features] block per attention head.
+    head_split_size: int = 0
+    head_split_param_names: list[str] = field(default_factory=list)
 
     @staticmethod
     def _get_default_load_balancing_config() -> LoadBalancingConfig:
@@ -1094,6 +1102,13 @@ def _get_default_load_balancing_config() -> LoadBalancingConfig:
         default_factory=_get_default_load_balancing_config
     )
 
+    def __post_init__(self) -> None:
+        super().__post_init__()
+        if self.head_split_size < 0:
+            raise ValueError(
+                f"Invalid {self.head_split_size=}. Must be >= 0 (0 disables head-split)."
+            )
+
 
 @dataclass(kw_only=True)
 class FSDPDistributedConfig(DistributedConfig):
@@ -1235,7 +1250,11 @@ class HybridShardDistributedConfig(FullyShardDistributedConfig, DDPDistributedCo
     device_mesh: DeviceMesh
 
 
-_ShampooPT2CompileConfigImpl: type[object] = make_dataclass(
+# Any, not type[object]: the fields are synthesized at import time from
+# torch.compile's signature, so a checker cannot resolve them. Any keeps the
+# subclass's attribute access permissive; type[object] instead made the class
+# statically unusable as a base.
+_ShampooPT2CompileConfigImpl: Any = make_dataclass(
     "_ShampooPT2CompileConfigImpl",
     [
         (name, param.annotation, param.default)
@@ -1246,9 +1265,7 @@ class HybridShardDistributedConfig(FullyShardDistributedConfig, DDPDistributedCo
 )
 
 
-class ShampooPT2CompileConfig(
-    _ShampooPT2CompileConfigImpl  # type: ignore
-):
+class ShampooPT2CompileConfig(_ShampooPT2CompileConfigImpl):
     """Configuration for Shampoo PT2 compilation.
 
     Enables Shampoo pytorch compilation with configure to speed up model training.
diff --git a/distributed_shampoo/tests/distributed_shampoo_test.py b/distributed_shampoo/tests/distributed_shampoo_test.py
index 2263b9c..4bd9c54 100644
--- a/distributed_shampoo/tests/distributed_shampoo_test.py
+++ b/distributed_shampoo/tests/distributed_shampoo_test.py
@@ -14,6 +14,7 @@
 import unittest
 from collections.abc import Callable
 from dataclasses import dataclass, field, replace
+from functools import partial
 from typing import Any, cast
 
 import torch
@@ -21,6 +22,7 @@
 from distributed_shampoo.preconditioner.matrix_functions_types import (
     DefaultNewtonSchulzOrthogonalizationConfig,
     EigenConfig,
+    NewtonSchulzRootInvConfig,
     OrthogonalizationConfig,
     PseudoInverseConfig,
 )
@@ -51,6 +53,11 @@
     TRAIN_MODE,
     WeightDecayType,
 )
+from distributed_shampoo.tests.shampoo_test_utils import (
+    compare_two_optimizers_on_weight_and_loss,
+    construct_training_problem,
+    train_model,
+)
 from distributed_shampoo.utils.shampoo_utils import pack_upper_triangular
 from torch import nn, Tensor
 from torch.testing._internal.common_utils import (
@@ -71,6 +78,114 @@ def _pack_if_enabled(
     )
 
 
+@instantiate_parametrized_tests
+class DistributedShampooNewtonSchulzTest(unittest.TestCase):
+    """End-to-end training with the Newton-Schulz inverse root instead of an eigendecomposition."""
+
+    @staticmethod
+    def _optim_factory(
+        parameters: Any,
+        preconditioner_config: PreconditionerConfig,
+    ) -> torch.optim.Optimizer:
+        return DistributedShampoo(
+            parameters,
+            lr=0.01,
+            betas=(0.9, 0.999),
+            epsilon=1e-8,
+            max_preconditioner_dim=5,
+            precondition_frequency=1,
+            start_preconditioning_step=1,
+            preconditioner_config=preconditioner_config,
+        )
+
+    def _assert_trains(self, preconditioner_config: PreconditionerConfig) -> None:
+        model, loss, data, target, _ = train_model(
+            optim_factory=partial(
+                DistributedShampooNewtonSchulzTest._optim_factory,
+                preconditioner_config=preconditioner_config,
+            ),
+            model_factory=partial(
+                construct_training_problem,
+                model_linear_layers_dims=(10, 5, 3),
+                model_dead_layers_dims=None,
+                fill=0.1,
+            ),
+            num_steps=20,
+        )
+        for name, parameter in model.named_parameters():
+            self.assertTrue(
+                torch.isfinite(parameter).all(),
+                msg=f"Non-finite values in {name} after training.",
+            )
+        # The problem targets zero, so a working preconditioner must drive the loss below its
+        # value at initialization.
+        self.assertLess(loss(model(data), target).item(), 0.1)
+
+    def test_training_with_newton_schulz(self) -> None:
+        self._assert_trains(
+            RootInvShampooPreconditionerConfig(
+                amortized_computation_config=NewtonSchulzRootInvConfig()
+            )
+        )
+
+    def test_training_with_coefficient_schedule(self) -> None:
+        """A per-iteration coefficient schedule, the form Polar Express supplies, trains end to end."""
+        self._assert_trains(
+            RootInvShampooPreconditionerConfig(
+                amortized_computation_config=NewtonSchulzRootInvConfig(
+                    coefficients=[[3.4445, -4.7750, 2.0315]] * 6
+                    + [[3.0, -16.0 / 5.0, 6.0 / 5.0]] * 14
+                )
+            )
+        )
+
+    def test_matches_eigendecomposition_on_full_rank_factor_matrices(self) -> None:
+        """Newton-Schulz agrees with the eigendecomposition once the factor matrices are full rank.
+
+        It does NOT agree while they are rank deficient, because relative_epsilon regularizes more
+        aggressively there than epsilon does on the eigendecomposition path.
+        """
+        compare_two_optimizers_on_weight_and_loss(
+            control_optim_factory=partial(
+                DistributedShampooNewtonSchulzTest._optim_factory,
+                preconditioner_config=DefaultShampooConfig,
+            ),
+            experimental_optim_factory=partial(
+                DistributedShampooNewtonSchulzTest._optim_factory,
+                preconditioner_config=RootInvShampooPreconditionerConfig(
+                    # Match the eigendecomposition path's regularization so the two are comparable.
+                    amortized_computation_config=NewtonSchulzRootInvConfig(
+                        relative_epsilon=0.0
+                    )
+                ),
+            ),
+            model_linear_layers_dims=(10, 10),
+            model_dead_layers_dims=None,
+            fill=0.1,
+            total_steps=5,
+            rtol=1e-2,
+            atol=1e-3,
+        )
+
+    def test_unsupported_root_fails_fast(self) -> None:
+        """An order-3 block asks for root 6, which Newton-Schulz cannot compute. This must raise at
+        optimizer construction rather than be swallowed as a per-factor-matrix warning during
+        training that silently reuses a stale preconditioner."""
+        model = nn.ParameterList([nn.Parameter(torch.randn(4, 4, 4))])
+        self.assertRaisesRegex(
+            ValueError,
+            re.escape(
+                "NewtonSchulzRootInvConfig only supports inverse roots that are powers of two, but "
+                "unsupported_roots=[6.0] were requested."
+            ),
+            DistributedShampooNewtonSchulzTest._optim_factory,
+            model.parameters(),
+            preconditioner_config=RootInvShampooPreconditionerConfig(
+                amortized_computation_config=NewtonSchulzRootInvConfig()
+            ),
+        )
+
+
 @instantiate_parametrized_tests
 class DistributedShampooInitTest(unittest.TestCase):
     def setUp(self) -> None:
diff --git a/distributed_shampoo/tests/shampoo_types_test.py b/distributed_shampoo/tests/shampoo_types_test.py
index c255e87..e204d69 100644
--- a/distributed_shampoo/tests/shampoo_types_test.py
+++ b/distributed_shampoo/tests/shampoo_types_test.py
@@ -9,6 +9,8 @@
 
 import re
 import unittest
+from dataclasses import asdict, fields
+from inspect import signature
 from typing import Any
 from unittest.mock import MagicMock
 
@@ -31,6 +33,7 @@
     HybridShardDistributedConfig,
     IterateAveragingConfig,
     RMSpropPreconditionerConfig,
+    ShampooPT2CompileConfig,
     SignDescentPreconditionerConfig,
 )
 from distributed_shampoo.utils.commons import get_all_non_abstract_subclasses
@@ -421,3 +424,24 @@ def test_illegal_num_sub_groups(self, num_sub_groups: int) -> None:
             device_mesh=MagicMock(),
             num_sub_groups=num_sub_groups,
         )
+
+
+class ShampooPT2CompileConfigTest(unittest.TestCase):
+    # The fields are synthesized at import time from torch.compile's signature, so no
+    # static checker can see them and nothing else in the suite asserts they exist.
+    def test_fields_match_torch_compile_signature(self) -> None:
+        self.assertEqual(
+            {field.name for field in fields(ShampooPT2CompileConfig())},
+            {name for name in signature(torch.compile).parameters if name != "model"},
+        )
+
+    def test_asdict_binds_to_torch_compile(self) -> None:
+        config = ShampooPT2CompileConfig(backend="eager", fullgraph=True)
+        kwargs = asdict(config)
+        self.assertEqual(kwargs["backend"], "eager")
+        self.assertTrue(kwargs["fullgraph"])
+        # Mirrors how distributed_shampoo.py splats the config into torch.compile.
+        signature(torch.compile).bind(torch.nn.Identity(), **kwargs)
+
+    def test_unknown_keyword_rejected(self) -> None:
+        self.assertRaises(TypeError, ShampooPT2CompileConfig, not_a_torch_compile_arg=1)
diff --git a/distributed_shampoo/utils/gpu_tests/shampoo_utils_test.py b/distributed_shampoo/utils/gpu_tests/shampoo_utils_test.py
index 994eb55..45b27ad 100644
--- a/distributed_shampoo/utils/gpu_tests/shampoo_utils_test.py
+++ b/distributed_shampoo/utils/gpu_tests/shampoo_utils_test.py
@@ -419,7 +419,7 @@ def test_param_recv_info_completeness(self) -> None:
         # Every param should have a valid recv info entry (offset >= 0)
         for param_idx in range(num_params):
             offset, chunk_size = ctx._param_recv_info[param_idx]
-            self.assertGreaterEqual(  # type: ignore
+            self.assertGreaterEqual(
                 offset, 0, f"Param {param_idx} has invalid recv offset {offset}"
             )
             self.assertGreaterEqual(
@@ -545,7 +545,7 @@ def test_gather_gradients_matches_full_tensor(self, num_params: int) -> None:
         # Verify correctness for assigned params
         for i in range(num_params):
             if i % self.world_size == rank:
-                self.assertIsNotNone(  # type: ignore
+                self.assertIsNotNone(
                     gathered_grads[i],
                     f"Assigned param {i} should have a gathered gradient",
                 )
@@ -556,7 +556,7 @@ def test_gather_gradients_matches_full_tensor(self, num_params: int) -> None:
                 )
             else:
                 # Unassigned params should be None
-                self.assertIsNone(  # type: ignore
+                self.assertIsNone(
                     gathered_grads[i],
                     f"Unassigned param {i} should be None on rank {rank}",
                 )
@@ -589,9 +589,7 @@ def test_gather_gradients_with_none_grads(self) -> None:
 
         # All should be None since no gradients were set
         for i, grad in enumerate(gathered_grads):
-            self.assertIsNone(  # type: ignore
-                grad, f"Param {i} has no grad, should be None"
-            )
+            self.assertIsNone(grad, f"Param {i} has no grad, should be None")
 
     @with_comms
     @skip_if_lt_x_gpu(4)
@@ -639,14 +637,14 @@ def test_gather_gradients_with_partial_none_grads(self) -> None:
         for i in range(len(shapes)):
             if i % self.world_size == rank:
                 if expected_full_grads[i] is not None:
-                    self.assertIsNotNone(gathered_grads[i])  # type: ignore
+                    self.assertIsNotNone(gathered_grads[i])
                     torch.testing.assert_close(
                         gathered_grads[i], expected_full_grads[i]
                     )
                 else:
-                    self.assertIsNone(gathered_grads[i])  # type: ignore
+                    self.assertIsNone(gathered_grads[i])
             else:
-                self.assertIsNone(gathered_grads[i])  # type: ignore
+                self.assertIsNone(gathered_grads[i])
 
     @with_comms
     @skip_if_lt_x_gpu(4)
@@ -673,7 +671,7 @@ def test_gather_gradients_preserves_shape(self) -> None:
 
         for i in range(len(shapes)):
             if i % self.world_size == rank:
-                self.assertIsNotNone(gathered_grads[i])  # type: ignore
+                self.assertIsNotNone(gathered_grads[i])
                 self.assertEqual(
                     gathered_grads[i].shape,  # type: ignore
                     torch.Size(shapes[i]),
@@ -722,11 +720,11 @@ def test_gather_gradients_multiple_calls(self) -> None:
         # Verify second call produces correct results (different from first)
         for i in range(len(shapes)):
             if i % self.world_size == rank:
-                self.assertIsNotNone(second_grads[i])  # type: ignore
+                self.assertIsNotNone(second_grads[i])
                 torch.testing.assert_close(second_grads[i], expected_second_grads[i])
                 # Verify second call differs from first
                 # (2x gradient vs 1x gradient for sum)
-                self.assertFalse(  # type: ignore
+                self.assertFalse(
                     torch.equal(first_grads[i], second_grads[i]),  # type: ignore
                     f"Second gather should differ from first for param {i}",
                 )
@@ -823,7 +821,7 @@ def test_gather_params_matches_full_tensor(self, num_params: int) -> None:
         # Verify correctness for assigned params
         for i in range(num_params):
             if i % self.world_size == rank:
-                self.assertIsNotNone(  # type: ignore
+                self.assertIsNotNone(
                     gathered_params[i],
                     f"Assigned param {i} should have a gathered value",
                 )
@@ -833,7 +831,7 @@ def test_gather_params_matches_full_tensor(self, num_params: int) -> None:
                     msg=f"Gathered param {i} does not match full_tensor()",
                 )
             else:
-                self.assertIsNone(  # type: ignore
+                self.assertIsNone(
                     gathered_params[i],
                     f"Unassigned param {i} should be None on rank {rank}",
                 )
diff --git a/distributed_shampoo/utils/optimizer_modules.py b/distributed_shampoo/utils/optimizer_modules.py
index 5de1ee9..40d33b9 100644
--- a/distributed_shampoo/utils/optimizer_modules.py
+++ b/distributed_shampoo/utils/optimizer_modules.py
@@ -87,6 +87,7 @@ def save_to_state_dict(
 
             for key, value in states:
                 if isinstance(value, torch.Tensor):
+                    # pyrefly: ignore [bad-argument-type]
                     destination[key] = value if keep_vars else value.detach()
                 elif isinstance(value, OptimizerModule):
                     destination[key] = {}
@@ -220,7 +221,6 @@ def load_from_new_state_to_old_state(
                 old_state = type(old_state)(
                     (
                         load_from_new_state_to_old_state(
-                            # pyrefly: ignore [bad-argument-type]
                             old_state=old_value,
                             # pyrefly: ignore [bad-index]
                             new_state=new_state[i],