equinox icon indicating copy to clipboard operation
equinox copied to clipboard

eqx.filter_jit use on methods

Open ak24watch opened this issue 2 months ago • 1 comments

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.

ak24watch avatar Oct 22 '25 06:10 ak24watch

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.

patrick-kidger avatar Oct 24 '25 22:10 patrick-kidger