feat: add scaled CANS inverse root utility

Signed-off-by: mkhona <mkhona@nvidia.com>

Author: mkhona-nvidia

emerging_optimizers/soap/cans_utils.py

@@ -0,0 +1,132 @@
+# SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+import math
+
+import torch
+from torch import Tensor
+
+from emerging_optimizers import utils
+from emerging_optimizers.utils import FP32MatmulPrecT
+
+
+__all__ = ["scaled_cans_coupled_ns"]
+
+_CANS_COEFFS = (
+    (5.182503604966906, -5.178098480082684),
+    (2.586120737395915, -0.6479542005271643),
+    (2.567364126726186, -0.6454968804392178),
+    (2.520560084348265, -0.6393528082067044),
+    (2.410759275435182, -0.6248683598710716),
+    (2.1883348130094173, -0.5952022073798908),
+    (1.8595760874873613, -0.5504490972723968),
+    (1.589020160467417, -0.5126569802066718),
+    (1.5051653981684994, -0.5007377068751799),
+    (1.5, -0.5),
+)
+_CANS_M = 16384.0
+_CANS_Z_MAX = (
+    1.0,
+    5.183,
+    13.403,
+    34.409,
+    86.731,
+    209.09,
+    457.55,
+    850.85,
+    1352.0,
+    2035.0,
+    3052.5,
+)
+_CANS_Y_MAX = (1.0, 1.297, 1.726, 1.473, 1.545, 1.415, 1.273, 1.148, 1.0, 1.0, 1.0)
+_CANS_P_MAX = (1.0, 3.98, 3.95, 3.88, 3.71, 3.32, 2.60, 1.73, 1.15, 1.008, 1.0)
+_CANS_S_Z = tuple(_CANS_M / value for value in _CANS_Z_MAX)
+_CANS_S_Y = tuple(_CANS_M / value for value in _CANS_Y_MAX)
+_CANS_S_P = tuple(_CANS_M / value for value in _CANS_P_MAX)
+_CANS_ALPHA_Y_0 = _CANS_COEFFS[0][1] * _CANS_S_Y[1] / (_CANS_S_Y[0] ** 2)
+_CANS_BETA_Y_0 = _CANS_COEFFS[0][0] * _CANS_S_Y[1] / _CANS_S_Y[0]
+_CANS_Z_SCALE_0 = _CANS_COEFFS[0][1] * _CANS_S_Z[1] / _CANS_S_Y[0]
+_CANS_Z_DIAG_ADD_0 = _CANS_COEFFS[0][0] * _CANS_S_Z[1]
+_CANS_ALPHA_P = tuple(_CANS_S_P[k] / (_CANS_S_Z[k] * _CANS_S_Y[k]) for k in range(1, len(_CANS_COEFFS)))
+_CANS_ALPHA_Y = tuple(
+    _CANS_COEFFS[k][1] * _CANS_S_Y[k + 1] / (_CANS_S_Y[k] * _CANS_S_P[k]) for k in range(1, len(_CANS_COEFFS))
+)
+_CANS_BETA_Y = tuple(_CANS_COEFFS[k][0] * _CANS_S_Y[k + 1] / _CANS_S_Y[k] for k in range(1, len(_CANS_COEFFS)))
+_CANS_ALPHA_Z_BASE = tuple(
+    _CANS_COEFFS[k][1] * _CANS_S_Z[k + 1] / (_CANS_S_Z[k] * _CANS_S_P[k]) for k in range(1, len(_CANS_COEFFS))
+)
+_CANS_BETA_Z_BASE = tuple(_CANS_COEFFS[k][0] * _CANS_S_Z[k + 1] / _CANS_S_Z[k] for k in range(1, len(_CANS_COEFFS)))
+_CANS_ALPHA_Z = _CANS_ALPHA_Z_BASE[:-1] + (_CANS_ALPHA_Z_BASE[-1] / _CANS_S_Z[-1],)
+_CANS_BETA_Z = _CANS_BETA_Z_BASE[:-1] + (_CANS_BETA_Z_BASE[-1] / _CANS_S_Z[-1],)
+
+
+def _estimate_max_eigenvalue(x: Tensor, eps: float) -> Tensor:
+    n = x.size(-1)
+    diag = x.diagonal(dim1=-2, dim2=-1)
+    mean = diag.sum(dim=-1) / n
+    sq_norm = torch.sum(x**2, dim=(-2, -1))
+    variance = torch.clamp((sq_norm / n) - (mean**2), min=0.0)
+    ws_bound = mean + torch.sqrt(variance) * math.sqrt(n - 1)
+
+    abs_x = torch.abs(x)
+    row_sum = torch.sum(abs_x, dim=-1).clamp_min_(eps)
+    minc_bound = torch.max(torch.einsum("...ij,...j->...i", abs_x, row_sum) / row_sum, dim=-1).values
+    return torch.minimum(ws_bound, minc_bound)
+
+
+def scaled_cans_coupled_ns(
+    x: Tensor,
+    eps: float = 1e-12,
+    fp32_matmul_prec: FP32MatmulPrecT = "highest",
+) -> Tensor:
+    """Compute inverse square roots with scaled coupled CANS Newton-Schulz.
+
+    Args:
+        x: A 2D symmetric positive-definite matrix or 3D batch of matrices.
+        eps: Lower bound used when normalizing the matrices.
+        fp32_matmul_prec: Precision used for FP32 matrix multiplications: ``"medium"`` for BF16,
+            ``"high"`` for TF32, or ``"highest"`` for FP32.
+
+    Returns:
+        The approximate inverse square root with the same shape as ``x``.
+    """
+    if x.dim() not in (2, 3) or x.shape[-2] != x.shape[-1]:
+        raise TypeError(f"x must be a square matrix or batch of square matrices, got shape {tuple(x.shape)}")
+    if not x.is_cuda:
+        raise TypeError("scaled_cans_coupled_ns only supports CUDA tensors")
+
+    with utils.fp32_matmul_precision(fp32_matmul_prec):
+        input_dtype = x.dtype
+        x = x.float()
+        max_eigval = (_estimate_max_eigenvalue(x, eps) * 1.01).clamp_min_(eps)
+        max_eigval = max_eigval[..., None, None]
+        cans_addmm = torch.addmm if x.dim() == 2 else torch.baddbmm
+        y = x * (_CANS_S_Y[0] / max_eigval)
+
+        z = y.mul(_CANS_Z_SCALE_0)
+        z.diagonal(dim1=-2, dim2=-1).add_(_CANS_Z_DIAG_ADD_0)
+        y = cans_addmm(y, y, y, beta=_CANS_BETA_Y_0, alpha=_CANS_ALPHA_Y_0)
+
+        for k in range(8):
+            p = cans_addmm(z, z, y, beta=0.0, alpha=_CANS_ALPHA_P[k])
+            y, z = (
+                cans_addmm(y, y, p, beta=_CANS_BETA_Y[k], alpha=_CANS_ALPHA_Y[k]),
+                cans_addmm(z, p, z, beta=_CANS_BETA_Z[k], alpha=_CANS_ALPHA_Z[k]),
+            )
+
+        p = cans_addmm(z, z, y, beta=0.0, alpha=_CANS_ALPHA_P[8])
+        inverse_root = cans_addmm(z, p, z, beta=_CANS_BETA_Z[8], alpha=_CANS_ALPHA_Z[8])
+        inverse_root.mul_(torch.rsqrt(max_eigval))
+        inverse_root = (inverse_root + inverse_root.mT) / 2.0
+        return inverse_root.to(input_dtype)

tests/test_cans_utils.py

@@ -0,0 +1,70 @@
+# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+import torch
+from absl import flags, logging
+from absl.testing import absltest, parameterized
+
+from emerging_optimizers.soap.cans_utils import scaled_cans_coupled_ns
+from emerging_optimizers.utils import FP32MatmulPrecT
+
+
+flags.DEFINE_enum("device", "cpu", ["cpu", "cuda"], "Device to run tests on")
+flags.DEFINE_integer("seed", None, "Random seed for reproducible tests")
+FLAGS = flags.FLAGS
+
+
+def setUpModule() -> None:
+    if FLAGS.seed is not None:
+        logging.info("Setting random seed to %d", FLAGS.seed)
+        torch.manual_seed(FLAGS.seed)
+        if torch.cuda.is_available():
+            torch.cuda.manual_seed_all(FLAGS.seed)
+
+
+class CANSUtilsTest(parameterized.TestCase):
+    @parameterized.parameters((4, 4), (2, 4, 4))  # type: ignore[misc]
+    def test_scaled_cans_supports_optional_batch_dimension(self, shape: tuple[int, ...]) -> None:
+        if FLAGS.device != "cuda":
+            self.skipTest("CANS requires CUDA")
+
+        x = torch.randn(*shape, device=FLAGS.device)
+        matrix = x @ x.mT + 0.1 * torch.eye(shape[-1], device=FLAGS.device)
+
+        inverse_root = scaled_cans_coupled_ns(matrix)
+
+        self.assertEqual(inverse_root.shape, matrix.shape)
+        self.assertTrue(torch.isfinite(inverse_root).all())
+
+    def test_scaled_cans_rejects_cpu_tensor(self) -> None:
+        with self.assertRaisesRegex(TypeError, "only supports CUDA"):
+            scaled_cans_coupled_ns(torch.eye(4))
+
+    @parameterized.parameters("medium", "high", "highest")  # type: ignore[misc]
+    def test_scaled_cans_supports_fp32_matmul_precision(self, fp32_matmul_prec: FP32MatmulPrecT) -> None:
+        if FLAGS.device != "cuda":
+            self.skipTest("CANS requires CUDA")
+
+        x = torch.randn(4, 4, device=FLAGS.device)
+        matrix = x @ x.T + 0.1 * torch.eye(4, device=FLAGS.device)
+        previous_precision = torch.get_float32_matmul_precision()
+
+        inverse_root = scaled_cans_coupled_ns(matrix, fp32_matmul_prec=fp32_matmul_prec)
+
+        self.assertTrue(torch.isfinite(inverse_root).all())
+        self.assertEqual(torch.get_float32_matmul_precision(), previous_precision)
+
+
+if __name__ == "__main__":
+    absltest.main()