How to keep hashable arguments static through a scan

[HenriLamarre]

Hi, disclaimer: I am new to JAX but it has been awesome so far! :D

I have a function ('jitfunc') that is jit-inherited and works properly. This function is being called in a while loop. The while loop itself is very slow whereas the function is fast. I would like to convert my while loop to either lax.while_loop, or lax.fori_loop but most likely lax.scan as it is differentiable and I the next step for this project is AI stuff.

So I have a maximal number of iterations itermax for my loop but the reason I use a while loop instead of a for loop is because sometimes, the jit-function can speed up the calculations and increase the iterand by more than 1.
In that scenario, I was planning to use scan, and have the function do nothing when it has reached itermax naturally (I think this is doable but I have no idea if it is optimal)

However, my issue comes from the fact that when I try running lax.scan, I get the error:
ValueError: Non-hashable static arguments are not supported, as this can lead to unexpected cache-misses. Static argument (index 6) of type for function in_loop is non-hashable.

From looking around online, it looks like using lax.scan makes my hashable static variables non-hashable? (Not sure about this one)
Is there a fix for what I am trying to do?

My code is not super readable so here is a minimal example that replicates my error:

from jax import jit
from jax import lax
from functools import partial


def loop_function(itermax):
    iter = 0
    nonstatic = 1
    static=2
    a, x = lax.scan(scanner, (iter, nonstatic, static), xs=None, length=itermax)

def scanner(a, x):
    iter = a[0]
    nonstatic = a[1]
    static = a[2]
    a, x = jitfunc(x, iter, nonstatic, static)
    return a, x

@partial(jit, static_argnums=(3,))
def jitfunc(x, iter, nonstatic, static):
    iter, nonstatic = process(iter, nonstatic, static)
    a = (iter, nonstatic, static)
    return a, x

@partial(jit, static_argnums=(2,))
def process(iter, nonstatic, static):
    nonstatic = nonstatic*static
    iter +=1+nonstatic%100
    return iter, nonstatic

if __name__=='__main__':
    loop_function(1000)

Thanks! :)

[jakevdp]

Thanks for the question! Any of the arguments that pass through the scan will be non-static. If you want arguments to be static, the best way to do so is probably via closure, for example:

def loop_function(itermax):
    iter = 0
    nonstatic = 1
    static = 2
    def scanner_closure(a, x):
        a, x = scanner((a[0], a[1], static), x)
        return (a[0], a[1]), x
    a, x = lax.scan(scanner_closure, (iter, nonstatic), xs=None, length=itermax)

Now static will be a static variable within the scanner function.

[erikrecio]

Hello! I'm facing a similar problem, I tried to define an argument as static, but it doesn't behave like one. This small code inspired from the example above gives me an error in the line if i in static:

import jax.numpy as jnp

def scanner(a, b, static, x):
    for i in range(20):
        if i in static:
            print(i, True)
        else:
            print(i, False)
    return (a, b), x

def loop_function(itermax):
    iter = 0
    nonstatic = jnp.array(list(range(10)))
    static = jnp.array(list(range(10)))

    def scanner_closure(a, x):
        a, x = scanner(a[0], a[1], static, x)
        a[0] += 1
        return a, x
    
    a, x = jax.lax.scan(scanner_closure, (iter, nonstatic), xs=None, length=itermax)

loop_function(2)

[jakevdp]

The issue is that your argument is an array. The array value is static, but within a JIT-compiled context, any operation you attempt to do on an array (static or not) will be traced and compiled. It sounds like you want this membership check to be evaluated at compile-time and if so you can do that this way:

def scanner(a, b, static, x):
    for i in range(20):
        with jax.ensure_compile_time_eval():
          if i in static:
              print(i, True)
          else:
              print(i, False)
    return (a, b), x

[erikrecio]

It works! Thanks for your quick answers :)

[patrick-kidger]

This is a very common issue. As a follow-up, you may like equinox.partition and equinox.combine to automatically split out the static vs non-static arguments. For example see how this is done around jit and grad in the documentation here.