equinox icon indicating copy to clipboard operation
equinox copied to clipboard

Globally set default filter functions?

Open LouisDesdoigts opened this issue 1 year ago • 0 comments

Hey Patrick, I was wondering if there was a way to set the default behavior for the filter functions in some global manner? I am currently building an fully differentiable optical modelling framework that makes heavy use of Equinox (which is awesome by the way!).

My current issue is that we want to use the enum package to set some properties (essentially tracking units) which I have been able to get working, but by default the filter_jit function uses the is_array_like filter function, resulting in ConcretizationTypeError, as the integer value from the enum.IntEnum is turned into a tracer.

This can be trivially fixed by passing the is_inexact_array filter function to the filter_jit decorator, however this means that we can not use jit functionality without explicitly passing this extra argument anytime we want to use it, Therefore for our use-case it would make sense to be able to set this globally, so that filter_jit can be called natively.

Is there currently any functionality for something like this, or any plans to implement something like this in the future? I imagine something like the jax.config could be possible, but I am mostly a high-level user so I'm not entirely sure.

Anyway any help would be greatly appreciated, and feel free to add our package to #75 - We have lost of awesome optics and astrophysics projects in the pipeline working on JWST, Hubble, photonic design.... lots of awesome stuff coming! Thanks again.

LouisDesdoigts avatar Aug 12 '22 01:08 LouisDesdoigts