Agent.infer_states() - Jax not working

[aandriella]

Hi,
I got this error when I try to run the testing_large_latent_space.ipynb and the building_up_agent_loop.ipynb. It seems related to the new optimisation that uses jax. Am I the only one having this issue?

Traceback (most recent call last): File "../pymdp/examples/testing_large_latent_space.py", line 196, in <module> agents = training_step(agents, env, batch_size, num_timesteps, lag=25) ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ File "../pymdp/examples/testing_large_latent_space.py", line 136, in training_step output, sequences = evolve_trials(agent, env, batch_size, num_timesteps) ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ File "../pymdp/examples/testing_large_latent_space.py", line 125, in evolve_trials last, sequences = lax.scan(step_fn, init, jnp.arange(num_timesteps)) ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ File "../pymdp/examples/testing_large_latent_space.py", line 104, in step_fn args, beliefs, outcomes, actions = vect_uas(carry['args'], keys, outcomes, actions) ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ File "../pymdp/examples/testing_large_latent_space.py", line 81, in update_agent_state beliefs = agent.infer_states(outcomes, actions, *args) ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ TypeError: Agent.infer_states() takes 3 positional arguments but 5 were given

[conorheins]

I suspect it's related to changes in the infer_states API introduced in this commit: 8ee1426

@dimarkov can you confirm this? If so, I can fix the usage in those notebooks to reflect the updated syntax