Tweaked example

patrick-kidger · Mar 27, 2022 · 3af0cde · 3af0cde
1 parent b65f303
commit 3af0cde
Showing 1 changed file with 4 additions and 2 deletions.
diff --git a/docs/api/stateful.md b/docs/api/stateful.md
@@ -18,10 +18,12 @@ import jax.lax as lax
 import jax.numpy as jnp
 
 index = eqx.experimental.StateIndex()
-eqx.experimental.set_state(index, jnp.array(0))
+init = jnp.array(0)
+eqx.experimental.set_state(index, init)
 
+@jax.jit
 def scan_fun(_, __):
-    val = eqx.experimental.get_state(index)
+    val = eqx.experimental.get_state(index, like=init)
     val = val + 1
     eqx.experimental.set_state(index, val)
     return None, val