Merge pull request #604 from BalzaniEdoardo:fix_prox_lasso

PiperOrigin-RevId: 745656867

Author: JAXopt authors

jaxopt/_src/prox.py

@@ -70,7 +70,7 @@ def prox_lasso(x: Any,
   if l1reg is None:
     l1reg = 1.0
 
-  if type(l1reg) == float:
+  if jnp.isscalar(l1reg):
     l1reg = tree_util.tree_map(lambda y: l1reg*jnp.ones_like(y), x)
 
   def fun(u, v): return jnp.sign(u) * jax.nn.relu(jnp.abs(u) - v * scaling)

tests/prox_test.py

@@ -87,6 +87,13 @@ def test_prox_lasso(self):
     got = prox.prox_lasso(x, alpha)
     self.assertArraysAllClose(jnp.array(expected), jnp.array(got))
 
+    # test that works when regularizer is float and prox jit compiled
+    got = jax.jit(prox.prox_lasso)(x, 0.5)
+    expected0 = [self._prox_l1(x[0][i], 0.5) for i in range(len(x[0]))]
+    expected1 = [self._prox_l1(x[1][i], 0.5) for i in range(len(x[0]))]
+    expected = (jnp.array(expected0), jnp.array(expected1))
+    self.assertArraysAllClose(jnp.array(expected), jnp.array(got))
+
   def _prox_enet(self, x, lam, gamma):
     return (1.0 / (1.0 + lam * gamma)) * self._prox_l1(x, lam)