Freezing parameters with partial function or static_argnames argument.

[gerdm]

Hi,

I'm trying to understand why is it possible to use the static_argnames argument in jax.jit in some occasions, but not in others. Further, why is it that when it's not possible to use static_argnames, I can still freeze the parameters using partial before jitting and this works just fine?

Setup

First, I create the dataset

import jax
from functools import partial
from string import ascii_uppercase
from jax.flatten_util import ravel_pytree

n_elements = 10
key = jax.random.PRNGKey(314159)
key_length, key_vals = jax.random.split(key)
keys_vals = jax.random.split(key_vals, n_elements)

letters = ascii_uppercase[:n_elements]
lengths = jax.random.choice(key_length, 20, (n_elements,))

elements = {
    letter: jax.random.normal(key, (length,))
    for letter, key, length in zip(letters, keys_vals, lengths)
}

Next, I specify the static parameters that will be passed to a function.

vals, r_fn = ravel_pytree(elements)

config = {
    "fn": r_fn,
    "more_stuff": ...
}

I'm interested in defining a function that jits over vals and has the elements config as static.

jax.jit and specify static_argnames

Using the following function definition, I define f that takes both an array of values and a static config element
.

@partial(jax.jit, static_argnames=("config",))
def f(vals, config):
    return config["fn"](vals)["B"].sum()

Then, calling f(vals, config) results in an error

ValueError                                Traceback (most recent call last)
Cell In [11], line 1
----> 1 f(vals, config)

ValueError: Non-hashable static arguments are not supported. An error occurred during a call to 'f' while trying to hash an object of type <class 'dict'>, {'fn': <function ravel_pytree.<locals>.<lambda> at 0x7f088607e3b0>, 'more_stuff': Ellipsis}. The error was:
TypeError: unhashable type: 'dict'

Static config before jitting

Jitting and passing static_argnames does not work in the previous example. However, if I explicitly specify that config is static before jitting using @partial(partial, config=config) —used in this ugly way as a way of example—, I'm able to run the jitted code.

@partial(jax.jit, static_argnames=("config",)) # equivalent to simply passing @jax.jit
@partial(partial, config=config)
def f(vals, config):
    return config["fn"](vals)["B"].sum()

Calling f(vals, config) run just fine

DeviceArray(-0.35970473, dtype=float32)

f
> <CompiledFunction of functools.partial(<function f at 0x7f0400b3fc70>, config={'fn': <function ravel_pytree.<locals>.<lambda> at 0x7f088607e3b0>, 'more_stuff': Ellipsis})>

---

I thought that passing static_argnames worked in a way like passing partial to specify static parameters and jitting after freezing. What gives?

[jakevdp]

static_argnames operates differently than closures; in particular, closed-over arrays will be treated as dynamic arguments, while closed-over values that are non-arrays will be treated as static.

You can see this in the jaxpr for closed-over functions:

@jax.jit
def f(x, y):
  return x + y

x = jnp.arange(4.0)
y = jnp.ones(4)

print(jax.make_jaxpr(f)(x, y))
# { lambda ; a:f32[4] b:f32[4]. let
#     c:f32[4] = xla_call[
#       call_jaxpr={ lambda ; d:f32[4] e:f32[4]. let f:f32[4] = add d e in (f,) }
#       name=f
#     ] a b
#   in (c,) }

print(jax.make_jaxpr(partial(f, y=y))(a))
# { lambda a:f32[4]; b:f32[4]. let
#     c:f32[4] = xla_call[
#       call_jaxpr={ lambda ; d:f32[4] e:f32[4]. let f:f32[4] = add d e in (f,) }
#       name=f
#     ] b a
#   in (c,) }

In both cases y is treated as a dynamic variable. The only difference is that in the first it is an explicit argument (appearing after the semicolon: lambda ; a:f32[4] b:f32[4]) while in the second it is an implicit argument (appearing before the semicolon: lambda a:f32[4]; b:f32[4]).

[gerdm]

Thanks @jakevdp. That's quite interesting!

Considering this, what's the recommended way to write a function that takes a static pytree and is jit-compatible? From my example above, I'm able to jit-compile and run if I pass the function explicitly.

@partial(jax.jit, static_argnames=("r_fn",)) # equivalent to simply passing @jax.jit
def f(vals, r_fn):
    return r_fn(vals)["B"].sum()

f(vals, r_fn) # Compiles and run

But in the code I'm working on, I have a small fixed number of variable functions are are set once at the beginning of the program. I was looking to define something like f(x, pytree_funcs) instead of f(x, fa, fb, fc). I'm assuming that pytree_funcs = {"fa": fa, "fb": fb, "fc": fc}.

[mattjj]

Further, why is it that when it's not possible to use static_argnames, I can still freeze the parameters using partial before jitting and this works just fine?

The practical difference to keep in mind is that static_argnums/static_argnames exists for compilation caching purposes. Let's say you have a function like this:

def f(i, x):
  if i == 0:
    ...
  elif i == 1:
    ...

We can't do f_jit = jit(f) because f_jit(0, 3.14) will complain that it can't specialize the Python control flow sufficiently.

We could do something like

def f_wrapper(i, x):
  return jit(partial(f, i))(x)

but then we won't get a cache hit if we evaluate f_wrapper(0, 3.14) followed by f_wrapper(0, 2.7).

If we instead wrote f_wrapper = jit(f, static_argnums=(0,)) then we would get a cache hit, based on comparing equality of the argument marked as static.

So, use static_argnums/static_argnames when:

1. you have a function that wouldn't otherwise be jittable due to Python control flow (or you have some other special reason to want to bake in several versions of the same constant into programs);
2. the argument in question is hashable, and;
3. you want more cache hits than you'd get by closure.

If you don't care about getting cache hits for equal values of the argument in question, you can just use closure and it'll act the same.