Neural networks potential example and optax import about jax-md HOT 2 CLOSED

---------------------------------------------------------------------------
TypeError                                 Traceback (most recent call last)
<ipython-input-25-43e0542c4175> in <module>()
     17 # Initialize the simulation
     18 
---> 19 state = init_fn(key, test_positions[0], Si_mass)
     20 nbrs = neighbor_fn(state.position)
     21 

12 frames
/usr/local/lib/python3.6/dist-packages/haiku/_src/transform.py in apply_fn(params, state, rng, *args, **kwargs)
    295         rng, err_msg=(APPLY_RNG_STATE_ERROR if state else APPLY_RNG_ERROR))
    296     with base.new_context(params=params, state=state, rng=rng) as ctx:
--> 297       out = f(*args, **kwargs)
    298     return out, ctx.collect_state()
    299 

TypeError: model() missing 1 required positional argument: 'neighbor'