From ca14973a1f35cf596f31ad1123cdd805d4de1333 Mon Sep 17 00:00:00 2001 From: Chen-Kuan Liao Date: Wed, 10 Jun 2026 19:15:12 -0500 Subject: [PATCH] Add vmap reset() regression test for nnx.metrics (#5483) Construct a MultiMetric(loss=Average) under nnx.vmap so its state carries a leading batch axis, call reset(), then run a vmapped update and assert compute() returns the expected per-batch means. Locks in the shape-preserving reset() behavior tracked in #5483. --- tests/nnx/metrics_test.py | 23 +++++++++++++++++++++++ 1 file changed, 23 insertions(+) diff --git a/tests/nnx/metrics_test.py b/tests/nnx/metrics_test.py index 15cf53a4e..0a8285f27 100644 --- a/tests/nnx/metrics_test.py +++ b/tests/nnx/metrics_test.py @@ -265,6 +265,29 @@ def test_average(self, with_mask, scalar_values): ) ) + def test_vmap_reset_preserves_shape(self): + n = 3 + + @nnx.vmap(in_axes=0, out_axes=0) + def init_metrics(seed): + del seed + return nnx.MultiMetric(loss=nnx.metrics.Average('loss')) + + metrics = init_metrics(jnp.arange(n)) + self.assertEqual(metrics.loss.total[...].shape, (n,)) + + metrics.reset() + self.assertEqual(metrics.loss.total[...].shape, (n,)) + + @nnx.vmap(in_axes=(0, 0), out_axes=None) + def do_update(m, value): + m.update(loss=value) + + do_update(metrics, jnp.arange(n, dtype=jnp.float32)) + np.testing.assert_array_equal( + metrics.compute()['loss'], jnp.arange(n, dtype=jnp.float32) + ) + if __name__ == '__main__': absltest.main()