Update model instantiators definitions test file in jax

tests/jax/test_jax_model_instantiators_definition.py

@@ -39,7 +39,7 @@ def test_parse_input(capsys):
             assert item not in output, f"'{item}' in '{output}'"
     # Mixed operation
     input = 'OBSERVATIONS["joint"] + concatenate([net * ACTIONS[:, -3:]])'
-    statement = 'inputs["states"]["joint"] + jnp.concatenate([net * inputs["taken_actions"][:, -3:]], axis=-1)'
+    statement = 'states["joint"] + jnp.concatenate([net * taken_actions[:, -3:]], axis=-1)'
     output = _parse_input(str(input))
     assert output.replace("'", '"') == statement, f"'{output}' != '{statement}'"