_NumpyroOptim.init emits jax type for iter count

numpyro/optim.py

@@ -41,7 +41,7 @@ def init(self, params: _Params) -> _IterOptState:
         :return: initial optimizer state.
         """
         opt_state = self.init_fn(params)
-        return 0, opt_state
+        return np.array(0), opt_state
 
     def update(self, g: _Params, state: _IterOptState) -> _IterOptState:
         """

test/test_optimizers.py

@@ -37,3 +37,36 @@ def test_optim_multi_params(optim_class, args):
         opt_state = step(opt_state, opt)
     for _, param in opt.get_params(opt_state).items():
         assert np.allclose(param, np.zeros(3))
+
+
+# note: this is somewhat of a bruteforce test. testing directly from
+# _NumpyroOptim would probably be better
+@pytest.mark.parametrize('optim_class, args', [
+    (optim.Adam, (1e-2,)),
+    (optim.ClippedAdam, (1e-2,)),
+    (optim.Adagrad, (1e-1,)),
+    (optim.Momentum, (1e-2, 0.5,)),
+    (optim.RMSProp, (1e-2, 0.95)),
+    (optim.RMSPropMomentum, (1e-4,)),
+    (optim.SGD, (1e-2,))
+])
+def test_numpyrooptim_no_double_jit(optim_class, args):
+
+    opt = optim_class(*args)
+    state = opt.init(np.zeros(10))
+
+    my_fn_calls = 0
+
+    @jit
+    def my_fn(state, g):
+        nonlocal my_fn_calls
+        my_fn_calls += 1
+
+        state = opt.update(g, state)
+        return state
+
+    state = my_fn(state, np.ones(10)*1.)
+    state = my_fn(state, np.ones(10)*2.)
+    state = my_fn(state, np.ones(10)*3.)
+
+    assert my_fn_calls == 1