Improve optax.linear_algebra.nnls tests

PiperOrigin-RevId: 727736758

Author: rdyro

optax/_src/linear_algebra.py

@@ -333,7 +333,7 @@ def f(x_p_c, _):
 
     return (xn, pn, cn), None
 
-  x = jnp.zeros(A.shape[1])
+  x = jnp.zeros_like(b, shape=b.shape[:-1] + A.shape[-1:])
   p = x
   c = 0.
 

optax/_src/linear_algebra_test.py

@@ -209,30 +209,43 @@ def _gen_symmetrix_matrix(dim, condition_number):
         # No guarantee of success after e >= 7
         pass
 
-  @parameterized.product(n=[24, 32, 64], d=[24, 32, 64], zero_A=[False, True])
-  def test_nnls(self, n, d, zero_A, atol=1e-3, iters=10**3):  # pylint: disable=invalid-name
+  @parameterized.product(n=[24, 32], d=[24, 32], zero_lhs=[False, True],
+                         seed=[0], dtype=[jnp.float32, jnp.bfloat16])
+  def test_nnls(self, n, d, zero_lhs, seed, dtype, atol=1e-5):
     """Test non-negative least squares solver."""
-    A = jnp.zeros((n, d)) if zero_A else np.random.normal(size=(n, d))  # pylint: disable=invalid-name
-    b = np.random.normal(size=(n,))
+    keys = jax.random.split(jax.random.key(seed), 2)
+    if zero_lhs:
+      if not (dtype == jnp.float32 and n == 32 and d == 32 and seed == 0):
+        self.skipTest('Only 1 test case for zero_lhs=True')
+      A = jnp.zeros((n, d), dtype=dtype)  # pylint: disable=invalid-name
+    else:
+      A = jax.random.normal(keys[0], (n, d), dtype=dtype)  # pylint: disable=invalid-name
+    b = jax.random.normal(keys[1], (n,), dtype=dtype)
 
-    x = linear_algebra.nnls(A, b, iters=iters)
+    x10 = linear_algebra.nnls(A, b, iters=10)
+    x100 = linear_algebra.nnls(A, b, iters=100)
+    x1000 = linear_algebra.nnls(A, b, iters=100000)
 
-    with self.subTest('x is non-negative'):
-      assert jnp.allclose(x.clip(max=0), 0, atol=atol)
+    with self.subTest('x has the correct dtype'):
+      self.assertEqual(x10.dtype, dtype)
+      self.assertEqual(x100.dtype, dtype)
+      self.assertEqual(x1000.dtype, dtype)
 
-    try:
-      xr, _ = scipy.optimize.nnls(A, b, maxiter=iters)
-    except RuntimeError:
-      return
+    with self.subTest('x is non-negative'):
+      assert jnp.allclose(x10.clip(max=0), 0, atol=atol)
+      assert jnp.allclose(x100.clip(max=0), 0, atol=atol)
+      assert jnp.allclose(x1000.clip(max=0), 0, atol=atol)
 
-    with self.subTest('xr is non-negative'):
-      assert jnp.allclose(xr.clip(max=0), 0, atol=atol)
+    # we skip comparison to scipy.optimize.nnls as convergence is flaky (by
+    # design, this is an iterative algorithm)
 
-    d = jnp.square(A @ x - b).sum()
-    dr = jnp.square(A @ xr - b).sum()
+    l10 = jnp.square(A @ x10 - b).sum()
+    l100 = jnp.square(A @ x100 - b).sum()
+    l1000 = jnp.square(A @ x1000 - b).sum()
 
-    with self.subTest('x is optimal'):
-      np.testing.assert_allclose(d, dr.clip(max=d), atol=atol)
+    with self.subTest('x is converging'):
+      jnp.allclose((l100 - l10).clip(max=0), 0, atol=atol)
+      jnp.allclose((l1000 - l100).clip(max=0), 0, atol=atol)
 
 
 if __name__ == '__main__':