aura-optax 0.1.0__tar.gz

aura_optax-0.1.0/.gitignore

@@ -0,0 +1,15 @@
+# Build artefacts
+build/
+dist/
+*.egg-info/
+
+# Python caches
+__pycache__/
+*.py[cod]
+.pytest_cache/
+.mypy_cache/
+.ruff_cache/
+
+# Virtual environments
+.venv/
+venv/

aura_optax-0.1.0/LICENSE

@@ -0,0 +1,21 @@
+MIT License
+
+Copyright (c) 2026 Enrico Ballini
+
+Permission is hereby granted, free of charge, to any person obtaining a copy
+of this software and associated documentation files (the "Software"), to deal
+in the Software without restriction, including without limitation the rights
+to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+copies of the Software, and to permit persons to whom the Software is
+furnished to do so, subject to the following conditions:
+
+The above copyright notice and this permission notice shall be included in all
+copies or substantial portions of the Software.
+
+THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+SOFTWARE.

aura_optax-0.1.0/PKG-INFO

@@ -0,0 +1,127 @@
+Metadata-Version: 2.5
+Name: aura-optax
+Version: 0.1.0
+Summary: AURA (Angular Update Rate Adaptation): a per-parameter step-size multiplier for Adam and Muon on complex-valued and real-valued parameters, as Optax gradient transformations
+Project-URL: Repository, https://github.com/enricoballini/aura-optax
+Project-URL: Issues, https://github.com/enricoballini/aura-optax/issues
+Author: Enrico Ballini
+License-Expression: MIT
+License-File: LICENSE
+Keywords: adam,complex-valued neural networks,jax,muon,optax,optimizer,step size
+Classifier: Development Status :: 4 - Beta
+Classifier: Intended Audience :: Science/Research
+Classifier: Programming Language :: Python :: 3
+Classifier: Programming Language :: Python :: 3 :: Only
+Classifier: Topic :: Scientific/Engineering :: Artificial Intelligence
+Requires-Python: >=3.10
+Requires-Dist: jax>=0.5.3
+Requires-Dist: optax>=0.2.7
+Provides-Extra: test
+Requires-Dist: numpy; extra == 'test'
+Requires-Dist: pytest>=7; extra == 'test'
+Description-Content-Type: text/markdown
+
+# aura-optax
+
+AURA (Angular Update Rate Adaptation) for [Optax](https://github.com/google-deepmind/optax).
+
+AURA scales the update direction of a base optimizer by a per-parameter multiplier γ, which increases when consecutive directions agree in the complex plane and decreases when they do not. It is designed for complex-valued neural networks and applies to real-valued parameters as well. The package provides three `optax.GradientTransformation`s:
+
+- `aura_optax.scale_by_aura`: AURA's multiplier alone, to be chained after the direction of any Optax optimizer;
+- `aura_optax.adam_aura`: the Adam direction, scaled by AURA's multiplier, with decoupled weight decay;
+- `aura_optax.muon_aura`: the Muon direction on weight matrices and the Adam direction on all other parameters, scaled by AURA's multiplier, with decoupled weight decay.
+
+## Installation
+
+```bash
+pip install aura-optax
+```
+
+The package requires Python ≥ 3.10, JAX and Optax ≥ 0.2.7. For GPU support, install the JAX build for your platform first, following the [JAX installation guide](https://docs.jax.dev/en/latest/installation.html).
+
+## Usage
+
+`scale_by_aura` is a gradient transformation that multiplies each update by γ. It is placed in an `optax.chain` after the transformation that produces the direction (`optax.scale_by_adam`, `optax.contrib.scale_by_muon`, ...) and before the learning rate:
+
+```python
+import optax
+
+import aura_optax
+
+optimizer = optax.chain(
+    optax.scale_by_adam(),  # the direction of any Optax optimizer
+    aura_optax.scale_by_aura(),  # AURA's multiplier
+    optax.add_decayed_weights(1e-4),  # optional decoupled weight decay
+    optax.scale_by_learning_rate(1e-2),
+)
+```
+
+The complete training loop below fits a complex-valued linear model by least squares with this optimizer:
+
+```python
+import jax
+import jax.numpy as jnp
+import optax
+
+import aura_optax
+
+key_x, key_w, key_b, key_init = jax.random.split(jax.random.key(0), 4)
+x = jax.random.normal(key_x, (256, 8), dtype=jnp.complex64)
+y = (
+    x @ jax.random.normal(key_w, (4, 8), dtype=jnp.complex64).T
+    + jax.random.normal(key_b, (4,), dtype=jnp.complex64)
+)
+params = {
+    "w": 0.1 * jax.random.normal(key_init, (4, 8), dtype=jnp.complex64),
+    "b": jnp.zeros(4, dtype=jnp.complex64),
+}
+
+
+def loss_fn(params):
+    residual = x @ params["w"].T + params["b"] - y
+    return jnp.mean(jnp.real(residual * jnp.conj(residual)))
+
+
+optimizer = optax.chain(
+    optax.scale_by_adam(),
+    aura_optax.scale_by_aura(),
+    optax.add_decayed_weights(1e-4),
+    optax.scale_by_learning_rate(1e-2),
+)
+opt_state = optimizer.init(params)
+
+
+@jax.jit
+def train_step(params, opt_state):
+    loss, grads = jax.value_and_grad(loss_fn)(params)
+    grads = jax.tree.map(jnp.conj, grads)  # see the note below
+    updates, opt_state = optimizer.update(grads, opt_state, params)
+    return optax.apply_updates(params, updates), opt_state, loss
+
+
+for step in range(1000):
+    params, opt_state, loss = train_step(params, opt_state)
+print(f"final loss: {loss:.3e}")
+```
+
+**Complex gradients.** For a real-valued loss L of complex parameters w = x + iy, `jax.grad` returns ∂L/∂x − i ∂L/∂y, which is the complex conjugate of the gradient g = ∂L/∂x + i ∂L/∂y whose negative is the steepest-descent direction. The gradients must therefore be conjugated before `optimizer.update`, as in the example above. The same conjugation is required by every Optax optimizer applied to complex parameters, and it leaves real-valued parameters unchanged.
+
+**Ready-made optimizers.** `adam_aura` is the chain above; `muon_aura` applies AURA to the Muon direction on the 2-D parameters and to the Adam direction on all others, with its own gate values. Both take the learning rate as a scalar or an Optax schedule:
+
+```python
+optimizer = aura_optax.adam_aura(learning_rate=1e-2)
+optimizer = aura_optax.muon_aura(learning_rate=1e-2)
+```
+
+**Monitoring the multiplier.** AURA's state is the entry of `opt_state` at the position of `scale_by_aura` in the chain: `opt_state[1]` in the chain above and in `adam_aura`, `opt_state[2]` in `muon_aura`:
+
+```python
+gamma = opt_state[1].multiplier  # same tree structure as params
+```
+
+## Method
+TODO
+
+## License
+
+MIT; see [LICENSE](LICENSE).

aura_optax-0.1.0/README.md

@@ -0,0 +1,104 @@
+# aura-optax
+
+AURA (Angular Update Rate Adaptation) for [Optax](https://github.com/google-deepmind/optax).
+
+AURA scales the update direction of a base optimizer by a per-parameter multiplier γ, which increases when consecutive directions agree in the complex plane and decreases when they do not. It is designed for complex-valued neural networks and applies to real-valued parameters as well. The package provides three `optax.GradientTransformation`s:
+
+- `aura_optax.scale_by_aura`: AURA's multiplier alone, to be chained after the direction of any Optax optimizer;
+- `aura_optax.adam_aura`: the Adam direction, scaled by AURA's multiplier, with decoupled weight decay;
+- `aura_optax.muon_aura`: the Muon direction on weight matrices and the Adam direction on all other parameters, scaled by AURA's multiplier, with decoupled weight decay.
+
+## Installation
+
+```bash
+pip install aura-optax
+```
+
+The package requires Python ≥ 3.10, JAX and Optax ≥ 0.2.7. For GPU support, install the JAX build for your platform first, following the [JAX installation guide](https://docs.jax.dev/en/latest/installation.html).
+
+## Usage
+
+`scale_by_aura` is a gradient transformation that multiplies each update by γ. It is placed in an `optax.chain` after the transformation that produces the direction (`optax.scale_by_adam`, `optax.contrib.scale_by_muon`, ...) and before the learning rate:
+
+```python
+import optax
+
+import aura_optax
+
+optimizer = optax.chain(
+    optax.scale_by_adam(),  # the direction of any Optax optimizer
+    aura_optax.scale_by_aura(),  # AURA's multiplier
+    optax.add_decayed_weights(1e-4),  # optional decoupled weight decay
+    optax.scale_by_learning_rate(1e-2),
+)
+```
+
+The complete training loop below fits a complex-valued linear model by least squares with this optimizer:
+
+```python
+import jax
+import jax.numpy as jnp
+import optax
+
+import aura_optax
+
+key_x, key_w, key_b, key_init = jax.random.split(jax.random.key(0), 4)
+x = jax.random.normal(key_x, (256, 8), dtype=jnp.complex64)
+y = (
+    x @ jax.random.normal(key_w, (4, 8), dtype=jnp.complex64).T
+    + jax.random.normal(key_b, (4,), dtype=jnp.complex64)
+)
+params = {
+    "w": 0.1 * jax.random.normal(key_init, (4, 8), dtype=jnp.complex64),
+    "b": jnp.zeros(4, dtype=jnp.complex64),
+}
+
+
+def loss_fn(params):
+    residual = x @ params["w"].T + params["b"] - y
+    return jnp.mean(jnp.real(residual * jnp.conj(residual)))
+
+
+optimizer = optax.chain(
+    optax.scale_by_adam(),
+    aura_optax.scale_by_aura(),
+    optax.add_decayed_weights(1e-4),
+    optax.scale_by_learning_rate(1e-2),
+)
+opt_state = optimizer.init(params)
+
+
+@jax.jit
+def train_step(params, opt_state):
+    loss, grads = jax.value_and_grad(loss_fn)(params)
+    grads = jax.tree.map(jnp.conj, grads)  # see the note below
+    updates, opt_state = optimizer.update(grads, opt_state, params)
+    return optax.apply_updates(params, updates), opt_state, loss
+
+
+for step in range(1000):
+    params, opt_state, loss = train_step(params, opt_state)
+print(f"final loss: {loss:.3e}")
+```
+
+**Complex gradients.** For a real-valued loss L of complex parameters w = x + iy, `jax.grad` returns ∂L/∂x − i ∂L/∂y, which is the complex conjugate of the gradient g = ∂L/∂x + i ∂L/∂y whose negative is the steepest-descent direction. The gradients must therefore be conjugated before `optimizer.update`, as in the example above. The same conjugation is required by every Optax optimizer applied to complex parameters, and it leaves real-valued parameters unchanged.
+
+**Ready-made optimizers.** `adam_aura` is the chain above; `muon_aura` applies AURA to the Muon direction on the 2-D parameters and to the Adam direction on all others, with its own gate values. Both take the learning rate as a scalar or an Optax schedule:
+
+```python
+optimizer = aura_optax.adam_aura(learning_rate=1e-2)
+optimizer = aura_optax.muon_aura(learning_rate=1e-2)
+```
+
+**Monitoring the multiplier.** AURA's state is the entry of `opt_state` at the position of `scale_by_aura` in the chain: `opt_state[1]` in the chain above and in `adam_aura`, `opt_state[2]` in `muon_aura`:
+
+```python
+gamma = opt_state[1].multiplier  # same tree structure as params
+```
+
+## Method
+TODO
+
+## License
+
+MIT; see [LICENSE](LICENSE).

aura_optax-0.1.0/aura_optax/__init__.py

@@ -0,0 +1,8 @@
+"""AURA (Angular Update Rate Adaptation) optimizers for Optax."""
+
+from aura_optax._alias import adam_aura, muon_aura
+from aura_optax._aura import AuraState, scale_by_aura
+
+__version__ = "0.1.0"
+
+__all__ = ["AuraState", "adam_aura", "muon_aura", "scale_by_aura", "__version__"]

aura_optax-0.1.0/aura_optax/_alias.py

@@ -0,0 +1,105 @@
+"""Adam-AURA and Muon-AURA optimizers."""
+
+from typing import Any, Callable
+
+import optax
+
+from aura_optax._aura import scale_by_aura
+from aura_optax._muon import muon_hybrid_direction
+
+
+def adam_aura(
+    learning_rate: optax.ScalarOrSchedule,
+    *,
+    b1: float = 0.9,
+    b2: float = 0.999,
+    eps: float = 1e-8,
+    beta_zeta: float = 0.95,
+    epsilon_e: float = 1e-6,
+    chi_alignment: float = 0.7,
+    chi_opposition: float = 0.4,
+    psi_alignment: float = 0.015,
+    psi_opposition: float = 0.3,
+    eta_minus: float = 0.99,
+    eta_plus: float = 1.01,
+    gamma_min: float = 1e-3,
+    gamma_max: float = 1e3,
+    gamma_init: float = 1.0,
+    weight_decay: float = 1e-4,
+    mask: Any | Callable[[optax.Params], Any] | None = None,
+) -> optax.GradientTransformation:
+    """Adam scaled by AURA's multiplier, with decoupled weight decay; AURA's state is state[1]."""
+    return optax.chain(
+        optax.scale_by_adam(b1=b1, b2=b2, eps=eps, nesterov=False),
+        scale_by_aura(
+            beta_zeta=beta_zeta,
+            epsilon_e=epsilon_e,
+            chi_alignment=chi_alignment,
+            chi_opposition=chi_opposition,
+            psi_alignment=psi_alignment,
+            psi_opposition=psi_opposition,
+            eta_minus=eta_minus,
+            eta_plus=eta_plus,
+            gamma_min=gamma_min,
+            gamma_max=gamma_max,
+            gamma_init=gamma_init,
+        ),
+        optax.add_decayed_weights(weight_decay, mask),
+        optax.scale_by_learning_rate(learning_rate),
+    )
+
+
+def muon_aura(
+    learning_rate: optax.ScalarOrSchedule,
+    *,
+    ns_steps: int = 5,
+    beta: float = 0.95,
+    nesterov: bool = True,
+    adam_b1: float = 0.9,
+    adam_b2: float = 0.999,
+    eps: float = 1e-8,
+    matrix_lr_scale: float = 10.0,
+    beta_zeta: float = 0.95,
+    epsilon_e: float = 1e-6,
+    chi_alignment: float = 0.75,
+    chi_opposition: float = 0.4,
+    psi_alignment: float = 0.01,
+    psi_opposition: float = 0.2,
+    eta_minus: float = 0.99,
+    eta_plus: float = 1.01,
+    gamma_min: float = 1e-3,
+    gamma_max: float = 1e3,
+    gamma_init: float = 1.0,
+    weight_decay: float = 1e-4,
+    mask: Any | Callable[[optax.Params], Any] | None = None,
+) -> optax.GradientTransformation:
+    """Muon (2-D parameters) or Adam (others) scaled by AURA's multiplier; AURA's state is state[2]."""
+    matrix_direction, bias_direction = muon_hybrid_direction(
+        ns_steps=ns_steps,
+        beta=beta,
+        nesterov=nesterov,
+        adam_b1=adam_b1,
+        adam_b2=adam_b2,
+        eps=eps,
+        matrix_lr_scale=matrix_lr_scale,
+    )
+    # The two direction stages stay flat (not a nested chain) so that AuraState is state[2].
+    return optax.chain(
+        matrix_direction,
+        bias_direction,
+        scale_by_aura(
+            beta_zeta=beta_zeta,
+            epsilon_e=epsilon_e,
+            chi_alignment=chi_alignment,
+            chi_opposition=chi_opposition,
+            psi_alignment=psi_alignment,
+            psi_opposition=psi_opposition,
+            eta_minus=eta_minus,
+            eta_plus=eta_plus,
+            gamma_min=gamma_min,
+            gamma_max=gamma_max,
+            gamma_init=gamma_init,
+        ),
+        optax.add_decayed_weights(weight_decay, mask),
+        optax.scale_by_learning_rate(learning_rate),
+    )

aura_optax-0.1.0/aura_optax/_aura.py

@@ -0,0 +1,110 @@
+"""AURA's per-parameter step-size multiplier as an Optax gradient transformation."""
+
+from typing import NamedTuple
+
+import jax
+import jax.numpy as jnp
+from jax.tree_util import tree_map
+import optax
+
+
+class AuraState(NamedTuple):
+    """State of scale_by_aura; chi + 1j * psi = zeta_ema / (1 - beta_zeta**count)."""
+
+    count: jax.Array
+    previous_direction: optax.Updates
+    zeta_ema: optax.Updates
+    multiplier: optax.Updates
+
+
+def scale_by_aura(
+    *,
+    beta_zeta: float = 0.95,
+    epsilon_e: float = 1e-6,
+    chi_alignment: float = 0.7,
+    chi_opposition: float = 0.4,
+    psi_alignment: float = 0.015,
+    psi_opposition: float = 0.3,
+    eta_minus: float = 0.99,
+    eta_plus: float = 1.01,
+    gamma_min: float = 1e-3,
+    gamma_max: float = 1e3,
+    gamma_init: float = 1.0,
+) -> optax.GradientTransformation:
+    """Scale each update by AURA's multiplier; chain after the direction, before the learning rate."""
+    if not 0.0 <= beta_zeta < 1.0:
+        raise ValueError("beta_zeta must be in [0, 1)")
+    if epsilon_e <= 0.0:
+        raise ValueError("epsilon_e must be positive")
+    if not -1.0 <= chi_opposition < chi_alignment <= 1.0:
+        raise ValueError(
+            "chi_opposition and chi_alignment must satisfy "
+            "-1 <= chi_opposition < chi_alignment <= 1"
+        )
+    if not 0.0 <= psi_alignment < psi_opposition <= 1.0:
+        raise ValueError(
+            "psi_alignment and psi_opposition must satisfy "
+            "0 <= psi_alignment < psi_opposition <= 1"
+        )
+    if not 0.0 < eta_minus < 1.0:
+        raise ValueError("eta_minus must be in (0, 1)")
+    if eta_plus <= 1.0:
+        raise ValueError("eta_plus must exceed 1")
+    if not 0.0 < gamma_min <= 1.0:
+        raise ValueError("gamma_min must be in (0, 1]")
+    if gamma_max < 1.0:
+        raise ValueError("gamma_max must be at least 1")
+    if not gamma_min <= gamma_init <= gamma_max:
+        raise ValueError("gamma_init must be in [gamma_min, gamma_max]")
+
+    def init_fn(params):
+        zeros = tree_map(jnp.zeros_like, params)
+        multiplier = tree_map(
+            lambda param: jnp.full_like(jnp.real(param), gamma_init), params
+        )
+        return AuraState(
+            count=jnp.zeros((), dtype=jnp.int32),
+            previous_direction=zeros,
+            zeta_ema=zeros,
+            multiplier=multiplier,
+        )
+
+    def update_fn(updates, state, params=None):
+        del params
+        count = state.count + 1
+        bias_correction = 1.0 - beta_zeta**count
+
+        def average_consistency(direction, previous_direction, zeta_ema_previous):
+            z = direction * jnp.conj(previous_direction)
+            energy = jnp.abs(direction) ** 2 + jnp.abs(previous_direction) ** 2
+            zeta = 2.0 * z / (energy + epsilon_e)
+            return beta_zeta * zeta_ema_previous + (1.0 - beta_zeta) * zeta
+
+        def next_multiplier(gamma, zeta_ema_value):
+            zeta_hat = zeta_ema_value / bias_correction
+            chi = jnp.real(zeta_hat)
+            psi = jnp.imag(zeta_hat)
+            shrink = jnp.logical_or(chi <= chi_opposition, jnp.abs(psi) >= psi_opposition)
+            grow = jnp.logical_and(
+                jnp.logical_not(shrink),
+                jnp.logical_and(chi >= chi_alignment, jnp.abs(psi) <= psi_alignment),
+            )
+            return jnp.where(
+                shrink,
+                jnp.maximum(eta_minus * gamma, gamma_min),
+                jnp.where(grow, jnp.minimum(eta_plus * gamma, gamma_max), gamma),
+            )
+
+        zeta_ema = tree_map(
+            average_consistency, updates, state.previous_direction, state.zeta_ema
+        )
+        multiplier = tree_map(next_multiplier, state.multiplier, zeta_ema)
+        scaled = tree_map(lambda direction, gamma: gamma * direction, updates, multiplier)
+        return scaled, AuraState(
+            count=count,
+            previous_direction=updates,
+            zeta_ema=zeta_ema,
+            multiplier=multiplier,
+        )
+
+    return optax.GradientTransformation(init_fn, update_fn)

aura_optax-0.1.0/aura_optax/_muon.py

@@ -0,0 +1,53 @@
+"""Muon-AURA's direction: Muon on 2-D parameters, Adam on all others."""
+
+import jax.numpy as jnp
+from jax.tree_util import tree_map
+import optax
+
+
+def muon_hybrid_direction(
+    *,
+    ns_steps: int,
+    beta: float,
+    nesterov: bool,
+    adam_b1: float,
+    adam_b2: float,
+    eps: float,
+    matrix_lr_scale: float,
+) -> tuple[optax.GradientTransformation, optax.GradientTransformation]:
+    """The masked Muon (2-D leaves) and Adam (other leaves) stages, returned unchained."""
+    is_matrix = lambda tree: tree_map(lambda leaf: leaf.ndim == 2, tree)
+    is_not_matrix = lambda tree: tree_map(lambda leaf: leaf.ndim != 2, tree)
+    # Masked-away leaves are childless MaskedNode sentinels, which tree_map skips.
+    matrix_dimension_numbers = lambda tree: tree_map(
+        lambda leaf: optax.contrib.MuonDimensionNumbers(), tree
+    )
+    matrix_direction = optax.masked(
+        optax.chain(
+            # Needs Optax >= 0.2.7: earlier scale_by_muon rescales by shape itself and uses x.T, not x.T.conj().
+            optax.contrib.scale_by_muon(
+                ns_steps=ns_steps,
+                beta=beta,
+                eps=eps,
+                nesterov=nesterov,
+                weight_dimension_numbers=matrix_dimension_numbers,
+            ),
+            # optax.contrib.muon's default width scaling, sqrt(max(1, fan_out / fan_in)).
+            optax.stateless(
+                lambda updates, params: tree_map(
+                    lambda leaf: jnp.sqrt(
+                        jnp.maximum(1.0, leaf.shape[1] / leaf.shape[0])
+                    )
+                    * leaf,
+                    updates,
+                )
+            ),
+            optax.scale(matrix_lr_scale),
+        ),
+        is_matrix,
+    )
+    bias_direction = optax.masked(
+        optax.scale_by_adam(b1=adam_b1, b2=adam_b2, eps=eps, nesterov=False),
+        is_not_matrix,
+    )
+    return matrix_direction, bias_direction

aura_optax-0.1.0/pyproject.toml

@@ -0,0 +1,44 @@
+[build-system]
+requires = ["hatchling>=1.27"]
+build-backend = "hatchling.build"
+
+[project]
+name = "aura-optax"
+dynamic = ["version"]
+description = "AURA (Angular Update Rate Adaptation): a per-parameter step-size multiplier for Adam and Muon on complex-valued and real-valued parameters, as Optax gradient transformations"
+readme = "README.md"
+license = "MIT"
+license-files = ["LICENSE"]
+requires-python = ">=3.10"
+authors = [{ name = "Enrico Ballini" }]
+keywords = ["jax", "optax", "optimizer", "complex-valued neural networks", "adam", "muon", "step size"]
+classifiers = [
+    "Development Status :: 4 - Beta",
+    "Intended Audience :: Science/Research",
+    "Programming Language :: Python :: 3",
+    "Programming Language :: Python :: 3 :: Only",
+    "Topic :: Scientific/Engineering :: Artificial Intelligence",
+]
+dependencies = [
+    "jax>=0.5.3",
+    "optax>=0.2.7",
+]
+
+[project.urls]
+Repository = "https://github.com/enricoballini/aura-optax"
+Issues = "https://github.com/enricoballini/aura-optax/issues"
+
+[project.optional-dependencies]
+test = ["pytest>=7", "numpy"]
+
+[tool.hatch.version]
+path = "aura_optax/__init__.py"
+
+[tool.hatch.build.targets.wheel]
+packages = ["aura_optax"]
+
+[tool.hatch.build.targets.sdist]
+include = ["aura_optax", "tests", "README.md", "LICENSE", "pyproject.toml"]
+
+[tool.pytest.ini_options]
+testpaths = ["tests"]

aura_optax-0.1.0/tests/test_aura.py

@@ -0,0 +1,197 @@
+""" """
+
+import cmath
+
+import jax
+import jax.numpy as jnp
+import numpy as np
+import optax
+import pytest
+
+import aura_optax
+
+ETA_MINUS = 0.99
+ETA_PLUS = 1.01
+
+
+def complex_array(rng, *shape):
+    return jnp.asarray(
+        rng.standard_normal(shape) + 1j * rng.standard_normal(shape),
+        dtype=jnp.complex64,
+    )
+
+
+def complex_params(seed=0):
+    rng = np.random.default_rng(seed)
+    return {
+        "w1": complex_array(rng, 8, 4),
+        "b1": complex_array(rng, 8),
+        "w2": complex_array(rng, 3, 8),
+        "b2": complex_array(rng, 3),
+    }
+
+
+def final_multiplier(directions):
+    """Multiplier after feeding directions to scale_by_aura with its defaults."""
+    transform = aura_optax.scale_by_aura()
+    state = transform.init(directions[0])
+    for direction in directions:
+        output, state = transform.update(direction, state)
+        np.testing.assert_allclose(
+            output, state.multiplier * direction, rtol=1e-6, equal_nan=False
+        )
+    return state.multiplier
+
+
+def test_first_step_decreases_multiplier():
+    # d_0 = 0 gives chi_1 = 0 <= chi_opposition.
+    direction = complex_array(np.random.default_rng(0), 5)
+    np.testing.assert_allclose(final_multiplier([direction]), ETA_MINUS, rtol=1e-6)
+
+
+@pytest.mark.parametrize("dtype", [jnp.complex64, jnp.float32])
+def test_constant_direction_increases_multiplier(dtype):
+    # chi_t = (1 - 0.95^(t-1)) / (1 - 0.95^t): decrease at t = 1, no change at t = 2, 3, increase after.
+    rng = np.random.default_rng(0)
+    if dtype == jnp.float32:
+        direction = jnp.asarray(rng.standard_normal(5), dtype=dtype)
+    else:
+        direction = complex_array(rng, 5)
+    multiplier = final_multiplier([direction] * 50)
+    assert multiplier.dtype == jnp.float32
+    np.testing.assert_allclose(multiplier, ETA_MINUS * ETA_PLUS**47, rtol=1e-4)
+
+
+def test_opposite_directions_decrease_multiplier():
+    direction = complex_array(np.random.default_rng(0), 5)
+    multiplier = final_multiplier([direction * (-1) ** t for t in range(20)])
+    np.testing.assert_allclose(multiplier, ETA_MINUS**20, rtol=1e-4)
+
+
+def test_rotation_decreases_multiplier_despite_alignment():
+    # chi stays near cos(0.5) = 0.88, but |psi| >= psi_opposition from t = 3; t = 2 changes nothing.
+    direction = complex_array(np.random.default_rng(0), 5)
+    directions = [direction * cmath.exp(0.5j * t) for t in range(30)]
+    np.testing.assert_allclose(final_multiplier(directions), ETA_MINUS**29, rtol=1e-4)
+
+
+def test_adam_aura_with_fixed_multiplier_equals_adamw():
+    params = complex_params()
+    rng = np.random.default_rng(1)
+    gradients = [
+        jax.tree.map(lambda p: complex_array(rng, *p.shape), params) for _ in range(10)
+    ]
+    aura = aura_optax.adam_aura(1e-2, gamma_min=1.0, gamma_max=1.0)
+    adamw = optax.adamw(1e-2, weight_decay=1e-4)
+    aura_params, aura_state = params, aura.init(params)
+    adamw_params, adamw_state = params, adamw.init(params)
+    for grads in gradients:
+        aura_updates, aura_state = aura.update(grads, aura_state, aura_params)
+        adamw_updates, adamw_state = adamw.update(grads, adamw_state, adamw_params)
+        aura_params = optax.apply_updates(aura_params, aura_updates)
+        adamw_params = optax.apply_updates(adamw_params, adamw_updates)
+    jax.tree.map(
+        lambda a, b: np.testing.assert_allclose(a, b, rtol=1e-6, atol=0, equal_nan=False),
+        aura_params,
+        adamw_params,
+    )
+
+
+@pytest.fixture
+def x64():
+    jax.config.update("jax_enable_x64", True)
+    yield
+    jax.config.update("jax_enable_x64", False)
+
+
+@pytest.mark.parametrize("seed", [0, 1])
+def test_muon_aura_is_phase_equivariant(seed, x64):
+    # For the matrices this holds only if Newton-Schulz uses the conjugate transpose.
+    # Runs in complex128: the Newton-Schulz steps amplify complex64 rounding to ~1e-1.
+    params = jax.tree.map(lambda p: p.astype(jnp.complex128), complex_params(seed))
+    rng = np.random.default_rng(seed + 10)
+    grads = jax.tree.map(
+        lambda p: complex_array(rng, *p.shape).astype(jnp.complex128), params
+    )
+    phase = cmath.exp(0.7j)
+    rotated = jax.tree.map(lambda g: phase * g, grads)
+    optimizer = aura_optax.muon_aura(1.0, weight_decay=0.0)
+    state = optimizer.init(params)
+    updates, _ = optimizer.update(grads, state, params)
+    rotated_updates, _ = optimizer.update(rotated, state, params)
+    jax.tree.map(
+        lambda u, r: np.testing.assert_allclose(
+            r, phase * u, rtol=1e-9, atol=1e-10, equal_nan=False
+        ),
+        updates,
+        rotated_updates,
+    )
+
+
+@pytest.mark.parametrize(
+    "make_optimizer", [aura_optax.adam_aura, aura_optax.muon_aura], ids=["adam", "muon"]
+)
+def test_training_reduces_loss(make_optimizer):
+    rng = np.random.default_rng(0)
+    x = complex_array(rng, 64, 4)
+    y = x @ complex_array(rng, 3, 4).T + complex_array(rng, 3)
+    params = {"w": 0.1 * complex_array(rng, 3, 4), "b": jnp.zeros(3, jnp.complex64)}
+
+    def loss_fn(params):
+        residual = x @ params["w"].T + params["b"] - y
+        return jnp.mean(jnp.real(residual * jnp.conj(residual)))
+
+    optimizer = make_optimizer(1e-2)
+
+    @jax.jit
+    def step(params, state):
+        grads = jax.tree.map(jnp.conj, jax.grad(loss_fn)(params))
+        updates, state = optimizer.update(grads, state, params)
+        return optax.apply_updates(params, updates), state
+
+    state = optimizer.init(params)
+    initial_loss = float(loss_fn(params))
+    for _ in range(500):
+        params, state = step(params, state)
+    assert float(loss_fn(params)) < 1e-2 * initial_loss
+
+
+def test_schedule_and_jit():
+    params = complex_params()
+    grads = complex_params(seed=5)
+    schedule = optax.linear_schedule(1e-2, 1e-3, transition_steps=10)
+    for optimizer in (aura_optax.adam_aura(schedule), aura_optax.muon_aura(schedule)):
+        state = optimizer.init(params)
+        update = jax.jit(optimizer.update)
+        for _ in range(3):
+            updates, state = update(grads, state, params)
+        for leaf, param in zip(jax.tree.leaves(updates), jax.tree.leaves(params)):
+            assert leaf.shape == param.shape and leaf.dtype == param.dtype
+            assert bool(jnp.all(jnp.isfinite(jnp.abs(leaf))))
+
+
+def test_aura_state_position_in_chain():
+    params = complex_params()
+    adam_state = aura_optax.adam_aura(1e-3).init(params)
+    muon_state = aura_optax.muon_aura(1e-3).init(params)
+    assert isinstance(adam_state[1], aura_optax.AuraState)
+    assert isinstance(muon_state[2], aura_optax.AuraState)
+
+
+@pytest.mark.parametrize(
+    "kwargs",
+    [
+        {"beta_zeta": 1.0},
+        {"epsilon_e": 0.0},
+        {"chi_opposition": 0.8, "chi_alignment": 0.7},
+        {"psi_alignment": 0.3, "psi_opposition": 0.3},
+        {"eta_minus": 1.0},
+        {"eta_plus": 1.0},
+        {"gamma_min": 0.0},
+        {"gamma_max": 0.5},
+        {"gamma_init": 2.0, "gamma_max": 1.5},
+    ],
+)
+def test_invalid_hyperparameters_raise(kwargs):
+    with pytest.raises(ValueError):
+        aura_optax.scale_by_aura(**kwargs)