How to tell the jit all my supporting data is static in a clean manner (i.e. subclassing, alternate data structure, etc)?

[isaacgerg]

I have a function I want to minimize which takes a vector input x and computes a scalar output s. The function needs an incredible amount of support data which is static. I save all the support data in a dictionary of arrays.

So, my function has the form
s = f(x, y)

where x is a vector and the free parameters, y is the support data stored in dict of arrays e.g. y = {'v1':np.array([...]), ... 'vn':np.array([...])}

I want to be able to do:
compiled_f = jax.jit(f, static_argnums= 1)

It seems like this would be a common pattern for optimization problems. My thought is that this is already solved and I doing something wrong.

[jakevdp]

Static arguments to JIT must be hashable, so dicts and arrays cannot be passed as static arguments. But if you want to pass arrays that will be treated as constants to your function, you can do so by closing over them. For example:

from functools import partial
compiled_f = jax.jit(partial(f, y=y))
result = compiled_f(x)

Does that work for your use-case?

[jakevdp]

That's similar, but not identical. In your version all leaves of the pytree y will be treated as dynamic. In my version, only arrays will be treated as dynamic, and any other values will be treated as static constants. One or the other may be more performant depending on the contents of y and what f does with it, but generally avoiding tracing of static arguments allows the compiler to specialize computations more and eke-out more performance. Though if your static values are changing often from call-to-call, it may trigger more re-compiles. Hope that helps!

[isaacgerg]

y is maid up mostly of array which are all constant (everything in y is constant). If i understand you correctly, all my arrays are dynamic which means they are not being treated as compile time constants. Is that correct?

Would I get better performance if I expanded the contents of y as new function parameters of f and then denoted them all static during the call to jit? There's a level of detail I think I am not quite grokking yet. (This is precisely what i am trying to make more compact as there are 20 array in y, which handling all those constant parameters becomes unwieldy in the code hence the dict package).

Having to recompile is not an issue. This is a one shot optimization problem for now.

[jakevdp]

Static arguments must be hashable, and arrays aren't hashable, so they can't normally be marked static.

[isaacgerg]

But should I make them static for speed? I can subclass the device arrays so that they are hashable.

[jakevdp]

I've not done any benchmarks regarding that, but I'd be surprised if marking array inputs as static would yield any performance improvements.

[isaacgerg]

From one of the maintainers, anecdotally it seems making these constant arrays static will not yield any benefits. Thus, using the closure in the first response is likely best.

From @jakevdp,

from functools import partial
compiled_f = jax.jit(partial(f, y=y))
result = compiled_f(x)