add AI written tests

Signed-off-by: Hao Wu <skyw@nvidia.com>

Author: skyw

tests/muown_reference.py

@@ -0,0 +1,143 @@
+# 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.
+
+# Reference Muown implementation, adapted from the authors' code for use as a test oracle:
+#   https://github.com/kcc-lion/muown/blob/main/optim/muown.py
+#   Lion et al., "Muown: Row-Norm Control for Muon Optimization", arXiv:2605.10797 (paper: CC BY 4.0).
+# The repository did not declare a code license at the time of copying.
+"""Reference Muown implementation used as a test oracle."""
+
+from typing import Callable
+
+import torch
+from torch import Tensor
+from torch.optim.optimizer import Optimizer
+
+
+def _wn_pre_ns(W: Tensor, g: Tensor, v_norm: Tensor, grad_W: Tensor) -> tuple[Tensor, Tensor, Tensor]:
+    """Reconstruct direction v from (W, g, v_norm) and split grad_W into (grad_g, grad_v)."""
+    u = W / g
+    v = u * v_norm
+    grad_g = (grad_W * u).sum(dim=1, keepdim=True)
+    grad_v = (g / v_norm) * (grad_W - u * grad_g)
+    return v, grad_g, grad_v
+
+
+def _wn_recompose(W: Tensor, g: Tensor, v_new: Tensor) -> Tensor:
+    """Write W = g * v_new / ||v_new||_row in place and return the new row norms."""
+    v_norm_new = v_new.norm(dim=1, keepdim=True)
+    W.copy_(g * (v_new / v_norm_new))
+    return v_norm_new
+
+
+class MuownReference(Optimizer):
+    """Single-process reference Muown with an injected orthogonalization callable."""
+
+    def __init__(
+        self,
+        params,
+        orthogonalize_fn: Callable[[Tensor], Tensor],
+        lr: float = 3e-4,
+        momentum: float = 0.95,
+        nesterov: bool = False,
+        betas: tuple[float, float] = (0.9, 0.95),
+        weight_decay: float = 0.0,
+        adam_eps: float = 1e-8,
+    ) -> None:
+        self._orthogonalize_fn = orthogonalize_fn
+        defaults = dict(
+            lr=lr,
+            momentum=momentum,
+            nesterov=nesterov,
+            betas=betas,
+            weight_decay=weight_decay,
+            adam_eps=adam_eps,
+        )
+        super().__init__(params, defaults)
+
+    def _init_state_2d(self, p: Tensor, state: dict) -> None:
+        w_norm = p.data.norm(dim=1, keepdim=True)
+        state["g"] = w_norm.clone()
+        state["v_norm"] = w_norm.clone()
+        state["m_v"] = torch.zeros_like(p.data)
+        state["m_g"] = torch.zeros_like(w_norm)
+        state["v_g"] = torch.zeros_like(w_norm)
+        state["step"] = 0
+
+    @torch.no_grad()
+    def step(self, closure=None):
+        loss = None
+        if closure is not None:
+            with torch.enable_grad():
+                loss = closure()
+
+        for group in self.param_groups:
+            lr = group["lr"]
+            momentum = group["momentum"]
+            nesterov = group["nesterov"]
+            betas = group["betas"]
+            weight_decay = group["weight_decay"]
+            adam_eps = group["adam_eps"]
+
+            for p in group["params"]:
+                if p.grad is None:
+                    continue
+
+                grad = p.grad
+                state = self.state[p]
+
+                if len(state) == 0:
+                    self._init_state_2d(p, state)
+
+                state["step"] += 1
+                step = state["step"]
+
+                g = state["g"]
+                v_norm = state["v_norm"]
+                m_v = state["m_v"]
+                m_g = state["m_g"]
+                v_g = state["v_g"]
+                if weight_decay != 0.0:
+                    W_old = p.data.clone()
+
+                # Fused: reconstruct v + compute weight norm gradients
+                v, grad_g, grad_v = _wn_pre_ns(p.data, g, v_norm, grad)
+
+                # Muon update on v: momentum + orthogonalization
+                m_v.mul_(momentum).add_(grad_v)
+                if nesterov:
+                    update = grad_v.add(m_v, alpha=momentum)
+                else:
+                    update = m_v.clone()
+
+                # Injected orthogonalization (folds in the 0.2 * sqrt(max(m, n)) scaling).
+                update = self._orthogonalize_fn(update)
+                v_new = v.add(update, alpha=-lr)
+
+                # Adam update on g (small [out_features, 1] vectors)
+                beta1, beta2 = betas
+                m_g.mul_(beta1).add_(grad_g, alpha=1 - beta1)
+                v_g.mul_(beta2).addcmul_(grad_g, grad_g, value=1 - beta2)
+                bc1 = 1 - beta1**step
+                bc2 = 1 - beta2**step
+                g.addcdiv_(m_g / bc1, (v_g / bc2).sqrt().add_(adam_eps), value=-lr)
+
+                # Fused: recompose W = g * v_new / ||v_new||, writes directly into p.data
+                state["v_norm"] = _wn_recompose(p.data, g, v_new)
+                if weight_decay != 0.0:
+                    p.data.add_(W_old, alpha=-lr * weight_decay)
+                    g.copy_(p.data.norm(dim=1, keepdim=True))
+
+        return loss

tests/test_muown.py

@@ -0,0 +1,156 @@
+# 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 muown_reference import MuownReference
+
+from emerging_optimizers.orthogonalized_optimizers.muown import Muown
+
+
+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)
+    # Use the most precise fp32 matmul path so the reference (which calls the optimizer's
+    # orthogonalization outside the fp32_matmul_precision context) computes bit-comparable results.
+    torch.set_float32_matmul_precision("highest")
+
+
+class MuownTest(parameterized.TestCase):
+    def setUp(self):
+        super().setUp()
+        self.device = FLAGS.device
+
+    @parameterized.product(shape=[(8, 16), (16, 8), (33, 65)])
+    def test_smoke(self, shape):
+        """A few steps run and keep the weight finite."""
+        p = torch.nn.Parameter(torch.randn(shape, device=self.device))
+        opt = Muown([p], lr=1e-2, weight_decay=0.01)
+        for _ in range(3):
+            p.grad = torch.randn_like(p)
+            opt.step()
+        self.assertTrue(torch.isfinite(p).all())
+
+    def test_raises_on_non_2d(self):
+        """Muown supports 2D parameters only."""
+        for shape in [(8,), (2, 3, 4)]:
+            p = torch.nn.Parameter(torch.randn(shape, device=self.device))
+            p.grad = torch.randn_like(p)
+            opt = Muown([p], lr=1e-2)
+            with self.assertRaises(TypeError):
+                opt.step()
+
+    def test_raises_on_closure(self):
+        """Closures are not supported."""
+        p = torch.nn.Parameter(torch.randn((8, 16), device=self.device))
+        p.grad = torch.randn_like(p)
+        opt = Muown([p], lr=1e-2)
+        with self.assertRaises(ValueError):
+            opt.step(lambda: 0.0)
+
+    @parameterized.product(shape=[(8, 16), (16, 8)], weight_decay=[0.0, 0.1])
+    def test_row_norm_equals_magnitude_state(self, shape, weight_decay):
+        """The reparameterization invariant ||W_row|| == g holds after each step."""
+        p = torch.nn.Parameter(torch.randn(shape, device=self.device))
+        opt = Muown([p], lr=1e-2, weight_decay=weight_decay)
+        for _ in range(4):
+            p.grad = torch.randn_like(p)
+            opt.step()
+            row_norm = p.detach().norm(dim=1, keepdim=True)
+            torch.testing.assert_close(
+                row_norm,
+                opt.state[p]["g"],
+                atol=1e-5,
+                rtol=1e-5,
+                msg=lambda m: f"Row norms of W must equal the magnitude state g.\n\n{m}",
+            )
+
+    def test_weight_decay_shrinks_magnitude(self):
+        """Decoupled weight decay shrinks the magnitude g relative to the no-decay run."""
+        p_wd = torch.nn.Parameter(torch.randn((16, 32), device=self.device))
+        p_no = torch.nn.Parameter(p_wd.detach().clone())
+        opt_wd = Muown([p_wd], lr=1e-2, weight_decay=0.1)
+        opt_no = Muown([p_no], lr=1e-2, weight_decay=0.0)
+        for _ in range(5):
+            grad = torch.randn_like(p_wd)
+            p_wd.grad = grad.clone()
+            p_no.grad = grad.clone()
+            opt_wd.step()
+            opt_no.step()
+        self.assertLess(opt_wd.state[p_wd]["g"].sum().item(), opt_no.state[p_no]["g"].sum().item())
+
+    @parameterized.product(shape=[(8, 16), (16, 8), (33, 65)], momentum=[0.0, 0.95])
+    def test_agrees_with_reference(self, shape, momentum):
+        """Muown matches the authors' reference implementation (no weight decay).
+
+        The reference uses classic (heavy-ball) momentum while Muown uses EMA momentum; the two differ by
+        a constant factor that the scale-invariant Newton-Schulz orthogonalization removes, so the updates
+        agree. Both use the same injected orthogonalization, so only float rounding (lerp vs mul/add,
+        Newton-Schulz normalization) separates them — hence a tight tolerance rather than bit-identity.
+        """
+        p = torch.nn.Parameter(torch.randn(shape, device=self.device))
+        p_ref = torch.nn.Parameter(p.detach().clone())
+
+        opt = Muown(
+            [p],
+            lr=1e-2,
+            momentum=momentum,
+            weight_decay=0.0,
+            extra_scale_factor=0.2,
+            fp32_matmul_prec="highest",
+        )
+        # Hold orthogonalization identical: feed the reference Muown's own scaled_orthogonalize_fn.
+        opt_ref = MuownReference(
+            [p_ref],
+            orthogonalize_fn=opt.scaled_orthogonalize_fn,
+            lr=1e-2,
+            momentum=momentum,
+            nesterov=False,
+            weight_decay=0.0,
+        )
+
+        for _ in range(5):
+            grad = torch.randn_like(p)
+            p.grad = grad.clone()
+            p_ref.grad = grad.clone()
+            opt.step()
+            opt_ref.step()
+
+        torch.testing.assert_close(
+            p.detach(),
+            p_ref.detach(),
+            atol=1e-5,
+            rtol=1e-4,
+            msg=lambda m: f"Muown weight diverged from the reference implementation.\n\n{m}",
+        )
+        torch.testing.assert_close(
+            opt.state[p]["g"],
+            opt_ref.state[p_ref]["g"],
+            atol=1e-5,
+            rtol=1e-4,
+            msg=lambda m: f"Muown magnitude g diverged from the reference implementation.\n\n{m}",
+        )
+
+
+if __name__ == "__main__":
+    absltest.main()