JIT a least squares loss function in Jax

Question

I have a simple loss function that looks like this

        def loss(r, x, y):
            resid = f(r, x) - y
            return jnp.mean(jnp.square(resid))

I would like to optimize over the parameter r and use some static parameters x and y to compute the residual. All parameters in question are DeviceArrays.

In order to JIT this, I tried doing the following

        @partial(jax.jit, static_argnums=(1, 2))
        def loss(r, x, y):
            resid = f(r, x) - y
            return jnp.mean(jnp.square(resid))

but I get this error

jax._src.traceback_util.UnfilteredStackTrace: ValueError: Non-hashable static arguments are not supported, as this can lead to unexpected cache-misses. Static argument (index 1) of type <class 'numpy.ndarray'> for function loss is non-hashable.

I understand that from #6233 that this is by design but I was wondering what the workaround here is, as this seems like a very common use case where you have some fixed (input, output) training data pairs and some free variable.

Thanks for any tips!

EDIT: this is the error I get when I just try to use jax.jit

jax._src.traceback_util.UnfilteredStackTrace: jax._src.errors.ConcretizationTypeError: Abstract tracer value encountered where concrete value is expected: Traced<ShapedArray(bool[])>with<DynamicJaxprTrace(level=0/1)>
The problem arose with the `bool` function. 
While tracing the function loss at /path/to/my/script:9 for jit, this concrete value was not available in Python because it depends on the value of the argument 'r'.`

jakevdp · Accepted Answer · 2021-10-16T14:30:06.320

It sounds like you're thinking of static arguments as "values that don't vary between computations". In JAX's JIT, static arguments can better be thought of as "hashable compile-time constants". In your case, you don't have hashable compile-time constants; you have arrays, so you can just JIT-compile with no static args:

@jit
def loss(r, x, y):
    resid = f(r, x) - y
    return jnp.mean(jnp.square(resid))

If you really want the JAX machinery to know that your arrays are constant, you can do so by passing them via a closure or a partial; for example:

from functools import partial

def loss(r, x, y):
    resid = f(r, x) - y
    return jnp.mean(jnp.square(resid))
loss = jit(partial(loss, x=x, y=y))

However, for the type of computation you are doing, where the constants are arrays operated on by JAX array functions, these two approaches lead to basically identical lowered XLA code, so you may as well use the simpler one.

Thanks for the detailed answer. However, I get an error when I try to run it with just `@jit`. (in the OP for better readability). Following the stack trace, it looks like there's something wrong with `f`, but I was wondering if I could get some more background on what this type of error means. Following the link it provides leads me to the static arguments annotation. — Carpetfizz, Oct 16 '21 at 21:17
Actually, I'll make a separate post about this since it's not relevant to the OP. Thanks for your help! — Carpetfizz, Oct 16 '21 at 21:23

JIT a least squares loss function in Jax

1 Answers1