Static JAX arrays in Equinox static=True fields can break `eqx.filter_jit` / cause recompilation errors

Problem

PR #345 puts a JAX array (formerly, what was called policies in the agent class) into an equinox static=True field. This can cause failures with equinox.filter_jit or other transforms that treat non-array parts as static arguments. Upon re-jit/reuse, Equinox/JAX sometimes need to compare/hash statics; JAX arrays don’t support that (array truth value is ambiguous; arrays aren’t hashable), which can crash or trigger pathological recompilation.

Why we don't currently see it:

Our tests/notebooks largely use jax.jit directly, or trace once; we don’t repeatedly call filter_jit’d functions with changing static fields.

Future risks.

Code like “environment loop wrapped in a single filter_jit’d function” (an example that @dimarkov raised to me) or any pattern that reuses a filter_jit’d function with an Agent argument will eventually trigger this if any static field contains a JAX array.

import equinox as eqx
import jax.numpy as jnp

from pymdp.agent import Agent


def build_agent():
    A = [jnp.array([[0.9, 0.1], [0.1, 0.9]], dtype=jnp.float32)]
    B = [jnp.stack([jnp.eye(2, dtype=jnp.float32), jnp.flipud(jnp.eye(2, dtype=jnp.float32))], axis=-1)]
    return Agent(A=A, B=B, policy_len=1, batch_size=1)


@eqx.filter_jit
def uses_agent(agent, x):
    # Access a value from the static policies field.
    return x + agent.policies.policy_arr[0, 0, 0]


def main():
    agent = build_agent()
    print("first", uses_agent(agent, 0.0))
    print("second", uses_agent(agent, 1.0))

    agent2 = build_agent()
    print("third", uses_agent(agent2, 2.0))


if __name__ == "__main__":
    main()

• With JAX_LOG_COMPILES=1, uses_agent recompiles on every call (even with the same agent), which indicates cache misses from static JAX arrays.
• A direct jax.jit(..., static_argnums=0) with a JAX array fails with ValueError Non-hashable static arguments..., matching the failure mode described above