test: update test effected by the optax==0.2.7 release (#2137)

* test: update test effected by the `optax==0.2.7` release

* test: remove test skip

Author: Qazalbash

test/test_optimizers.py

@@ -100,7 +100,6 @@ def test_optim_multi_params(optim_class, args, kwargs, uses_value_arg):
     + optax_optimizers,
 )
 @pytest.mark.filterwarnings("ignore:.*tree_multimap:FutureWarning")
-@pytest.mark.skip(reason="Failing on optax==0.2.7. Halting till #2137 is merged.")
 def test_numpyrooptim_no_double_jit(optim_class, args, kwargs, uses_value_arg):
     opt = optim_class(*args, **kwargs)
     if not isinstance(opt, optim._NumPyroOptim):
@@ -124,8 +123,4 @@ def my_fn(state, g):
     state = my_fn(state, jnp.ones(10) * 2.0)
     state = my_fn(state, jnp.ones(10) * 3.0)
 
-    if uses_value_arg:
-        # Dtype is different on the first call vs the rest of the calls
-        assert my_fn_calls == 2
-    else:
-        assert my_fn_calls == 1
+    assert my_fn_calls == 1