bugfix relating to RNNO and celltype=lstm

src/ring/ml/rnno_v1.py

@@ -35,10 +35,11 @@ def forward_fn(X):
         assert X.shape[-2] == 1
 
         for i, n_units in enumerate(rnn_layers):
-            n_units = _factor * n_units
-            state = hk.get_state(f"rnn_{i}", shape=[1, n_units], init=jnp.zeros)
-            X, state = hk.dynamic_unroll(_cell(n_units), X, state)
-            hk.set_state(f"rnn_{i}", state)
+            state = hk.get_state(
+                f"rnn_{i}", shape=[1, n_units * _factor], init=jnp.zeros
+            )
+            X, state = hk.dynamic_unroll(_cell(n_units), X[..., 0, :], state[0])
+            hk.set_state(f"rnn_{i}", state[None])
 
             if layernorm:
                 X = hk.LayerNorm(axis=-1, create_scale=False, create_offset=False)(X)