Handling large arrays under jit

[renecotyfanboy]

Hey everyone, I am facing issues when it comes to interpolating large datasets. I end up in situations where I either need to wait for constant folding (I think this issue is even worse due to jax-ml/jax#23822) or use memory-mapping in a pure_callback which I find kinda inefficient. Any extra ideas to handle this ?