|
| 1 | +# Copyright 2025 BrainX Ecosystem Limited. All Rights Reserved. |
| 2 | +# |
| 3 | +# Licensed under the Apache License, Version 2.0 (the "License"); |
| 4 | +# you may not use this file except in compliance with the License. |
| 5 | +# You may obtain a copy of the License at |
| 6 | +# |
| 7 | +# http://www.apache.org/licenses/LICENSE-2.0 |
| 8 | +# |
| 9 | +# Unless required by applicable law or agreed to in writing, software |
| 10 | +# distributed under the License is distributed on an "AS IS" BASIS, |
| 11 | +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. |
| 12 | +# See the License for the specific language governing permissions and |
| 13 | +# limitations under the License. |
| 14 | +# ============================================================================== |
| 15 | +"""Regression tests for the 2026-06-19 optim/losses audit. |
| 16 | +
|
| 17 | +These pin the bug fixes for: |
| 18 | +
|
| 19 | +- P1-C1/P1-C2: ``Adan.update`` crashed on every call (``jax.lax.cond`` operand |
| 20 | + mis-binding) and its step counter was frozen at 0 (bias correction / Nesterov |
| 21 | + term disabled). |
| 22 | +- P1-H1: ``SM3`` raised ``KeyError`` for scalar (0-dim) trainable variables. |
| 23 | +""" |
| 24 | + |
| 25 | +import numpy as np |
| 26 | + |
| 27 | +import brainpy.math as bm |
| 28 | +from brainpy.optim import optimizer as O |
| 29 | + |
| 30 | + |
| 31 | +def _vec_var(values): |
| 32 | + return bm.Variable(bm.asarray(np.asarray(values, dtype=np.float32))) |
| 33 | + |
| 34 | + |
| 35 | +def _scalar_var(value=2.0): |
| 36 | + return bm.Variable(bm.asarray(np.asarray(value, dtype=np.float32))) |
| 37 | + |
| 38 | + |
| 39 | +def _train(opt, var_dict, grad_fn, steps=5): |
| 40 | + for _ in range(steps): |
| 41 | + grads = {k: grad_fn(bm.as_jax(v.value)) for k, v in var_dict.items()} |
| 42 | + opt.update(grads) |
| 43 | + return {k: np.asarray(bm.as_jax(v.value)) for k, v in var_dict.items()} |
| 44 | + |
| 45 | + |
| 46 | +# --------------------------------------------------------------------------- |
| 47 | +# Adan (P1-C1, P1-C2) |
| 48 | +# --------------------------------------------------------------------------- |
| 49 | +class TestAdanFixed: |
| 50 | + def test_adan_runs_and_updates(self): |
| 51 | + # Previously raised TypeError on the very first update. |
| 52 | + v = _vec_var([2.0, -3.0]) |
| 53 | + opt = O.Adan(lr=1e-2, train_vars={'w': v}) |
| 54 | + out = _train(opt, {'w': v}, lambda val: val, steps=5) |
| 55 | + assert np.all(np.isfinite(out['w'])) |
| 56 | + # gradient = value -> the parameter should move toward zero. |
| 57 | + assert np.all(np.abs(out['w']) < np.array([2.0, 3.0])) |
| 58 | + |
| 59 | + def test_adan_no_prox_runs(self): |
| 60 | + v = _vec_var([2.0, -3.0]) |
| 61 | + opt = O.Adan(lr=1e-2, train_vars={'w': v}, no_prox=True) |
| 62 | + out = _train(opt, {'w': v}, lambda val: val, steps=5) |
| 63 | + assert np.all(np.isfinite(out['w'])) |
| 64 | + |
| 65 | + def test_adan_step_counter_advances(self): |
| 66 | + v = _vec_var([1.0]) |
| 67 | + opt = O.Adan(lr=1e-2, train_vars={'w': v}) |
| 68 | + for _ in range(4): |
| 69 | + opt.update({'w': bm.as_jax(v.value)}) |
| 70 | + assert int(bm.as_jax(opt.step.value)) == 4 |
| 71 | + |
| 72 | + def test_adan_first_step_diff_is_zero(self): |
| 73 | + # On the very first update the gradient difference (g - g_prev) must be |
| 74 | + # treated as 0, so the exp_avg_diff (``_v``) accumulator stays 0 after |
| 75 | + # one step regardless of the gradient magnitude. |
| 76 | + v = _vec_var([5.0]) |
| 77 | + opt = O.Adan(lr=1e-2, train_vars={'w': v}) |
| 78 | + opt.update({'w': bm.as_jax(v.value)}) |
| 79 | + v_diff = np.asarray(bm.as_jax(opt.implicit_vars['w_v'].value)) |
| 80 | + assert np.allclose(v_diff, 0.0) |
| 81 | + |
| 82 | + def test_adan_nesterov_term_active_after_two_steps(self): |
| 83 | + # With the step counter frozen the diff term was permanently 0; here a |
| 84 | + # changing gradient must produce a non-zero exp_avg_diff after step 2. |
| 85 | + v = _vec_var([1.0]) |
| 86 | + opt = O.Adan(lr=1e-2, train_vars={'w': v}) |
| 87 | + opt.update({'w': np.asarray([1.0], dtype=np.float32)}) |
| 88 | + opt.update({'w': np.asarray([3.0], dtype=np.float32)}) # gradient changed |
| 89 | + v_diff = np.asarray(bm.as_jax(opt.implicit_vars['w_v'].value)) |
| 90 | + assert not np.allclose(v_diff, 0.0) |
| 91 | + |
| 92 | + |
| 93 | +# --------------------------------------------------------------------------- |
| 94 | +# SM3 (P1-H1) |
| 95 | +# --------------------------------------------------------------------------- |
| 96 | +class TestSM3Fixed: |
| 97 | + def test_sm3_scalar_var_runs(self): |
| 98 | + # Previously raised KeyError('w_m0') for a 0-dim variable. |
| 99 | + v = _scalar_var(2.0) |
| 100 | + opt = O.SM3(lr=0.1, train_vars={'w': v}) |
| 101 | + out = _train(opt, {'w': v}, lambda val: np.ones_like(val), steps=4) |
| 102 | + assert np.all(np.isfinite(out['w'])) |
| 103 | + # gradient is +1 each step -> scalar parameter must decrease. |
| 104 | + assert float(out['w']) < 2.0 |
| 105 | + |
| 106 | + def test_sm3_scalar_matches_adagrad_like_step(self): |
| 107 | + # For a scalar with constant gradient g=1, SM3 reduces to an Adagrad-like |
| 108 | + # update: cache accumulates g^2, step = lr * g / sqrt(cache + eps). |
| 109 | + v = _scalar_var(0.0) |
| 110 | + opt = O.SM3(lr=0.1, train_vars={'w': v}, eps=1e-30) |
| 111 | + opt.update({'w': np.asarray(1.0, dtype=np.float32)}) |
| 112 | + # after one step cache = 1, update = 0.1 * 1 / sqrt(1) = 0.1 |
| 113 | + assert float(bm.as_jax(v.value)) == np.float32(-0.1) |
| 114 | + |
| 115 | + def test_sm3_scalar_still_works_with_momentum(self): |
| 116 | + v = _scalar_var(2.0) |
| 117 | + opt = O.SM3(lr=0.1, train_vars={'w': v}, momentum=0.5, beta=0.5) |
| 118 | + out = _train(opt, {'w': v}, lambda val: np.ones_like(val), steps=3) |
| 119 | + assert np.all(np.isfinite(out['w'])) |
0 commit comments