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 :)

[yuichiro-student]

Can we make the closer with lambda, right?
Could you explain the difference between creating a new inner function using def and lambda?

When we change the closure, do these functions get recompiled, or cause slow down?

[jakevdp]

From the JAX perspective, there is no difference between defining the inner function with def vs. with lambda. From the Python perspective, one function is anonymous and one is not, but JAX will treat them identically.

[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.

[adonath]

For a native JAX solution I think handling the carry of a scanned function as namedtuple (in general) or dataclass (with static members) is the best idea. In your example it would look like:

import jax
from functools import partial
from dataclasses import dataclass, field


@jax.tree_util.register_dataclass
@dataclass
class ScanCarry:
    """Named scan args"""
    iter: jax.Array
    nonstatic: jax.Array
    static: int = field(metadata={"static": True})

def loop_function(itermax):
    carry = ScanCarry(
        iter=0,
        nonstatic=1,
        static=2,
    )
    return jax.lax.scan(body, carry, xs=None, length=itermax)

def body(a, x):
    nonstatic = a.nonstatic * a.static
    iter = a.iter + 1 + nonstatic % 100
    return ScanCarry(iter, nonstatic, a.static), x


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

I think it improves readability and gets rid of to many partial() and jit, because scanned functions are always jitted.