change a lot

examples/state_test.py

@@ -1,6 +1,8 @@
 import jax
+from jax import numpy as jnp
 from algorithm.state import State
 
+
 @jax.jit
 def func(state: State, a):
     return state.update(a=a)
@@ -9,6 +11,5 @@ def func(state: State, a):
 state = State(c=1, b=2)
 print(state)
 
-state = func(state, 1111111)
-
-print(state)
+vmap_func = jax.vmap(func, in_axes=(None, 0))
+print(vmap_func(state, jnp.array([1, 2, 3])))