equinox
equinox copied to clipboard
eqx.filter_jit use on methods
how eqx.filter_jit works on methods , is it same as jax.jit with static argnums on self. can anyone give explain it and give examples of usages.
It should work on methods the same way it works on regular functions, and self isn't privileged. Thus, any parts of the pytree structure of self that are arrays are traced, and any parts that are non-arrays are static.