Caching of static arguments

[cisprague]

I was searching the internet/documentation to see if jax.jit retains a cache for previously used static arguments to a function. I couldn't seem to find any information, so I did a quick test. It seems that jax.jit does retain the cache.

import jax, jax.numpy as jnp
from functools import partial
x = jax.random.normal(jax.random.PRNGKey(0))

class Function:
    @staticmethod
    @partial(jax.jit, static_argnums=(1,))
    def __call__(x, y):
        jnp.exp(x) if y else jnp.exp(x)
f = Function()

%timeit -n 100 -r 100 f(x, True)
%timeit -n 100 -r 100 f(x, False)
%timeit -n 100 -r 100 f(x, True)
%timeit -n 100 -r 100 f(x, False)

The slowest run took 8.54 times longer than the fastest. This could mean that an intermediate result is being cached.
12.3 µs ± 6.15 µs per loop (mean ± std. dev. of 100 runs, 100 loops each)
The slowest run took 9.37 times longer than the fastest. This could mean that an intermediate result is being cached.
14.1 µs ± 7.08 µs per loop (mean ± std. dev. of 100 runs, 100 loops each)
13.7 µs ± 3.89 µs per loop (mean ± std. dev. of 100 runs, 100 loops each)
13.4 µs ± 3.95 µs per loop (mean ± std. dev. of 100 runs, 100 loops each)

[jakevdp]

Yes, jit has a runtime cache so that subsequent calls to a JIT-compiled function will hit the cache rather than recompile. The cache key is constructed from the shape and dtype of input arrays, the hash of static arguments, and global configuration values that may affect execution. The details are internal, but there are a few APIs that let you inspect the cache; for example:

import jax

def f(x, flag):
  return x if flag else x + 1

f_jit = jax.jit(f, static_argnums=1)

print(f_jit._cache_size())  # 0

f_jit(1.0, True)
print(f_jit._cache_size())  # 1

f_jit(1.0, False)
print(f_jit._cache_size())  # 2

# re-wrapped function hits the same cache
f_jit_2 = jax.jit(f, static_argnums=1)
print(f_jit_2._cache_size())  # 2

# cache hits don't increase the cache size
f_jit(100.0, True)
f_jit(100.0, False)
print(f_jit._cache_size())  # 2

[cisprague]

Thanks for the response. I would like recompilation only to happen when an attribute relevant to the method changes. And I would like to do this without having to define anything outside of the class. In my case, I want to define a JITed parent class method that uses a child method, where the child method itself dynamically and statically makes use of some combination of its own attributes and its parent's attributes. Shouldn't this be possible by just comparing the method signature to the set of statically-treated attributes?

[jakevdp]

Shouldn't this be possible by just comparing the method signature to the set of statically-treated attributes?

I don't think this is possible. Say you have a function that looks like this:

def f(x, condition1, condition2):
  return x + 1 if condition1 else x + 2

How would you propose determining that the value of condition2 does not affect the output of the function, without executing it to see whether the results change?

[cisprague]

@jakevdp thanks for your help. It seems I'm able to get the behaviour I want in the following way.

@dataclass
class Parent:

    dp: float # dynamic parameter
    sp: bool # static parameter

    def kwargs(self):
        return (dict(dp=self.dp), dict(sp=self.sp))

    def __call__(self, x: float) -> float:
        # API method
        dynamic, static = self.kwargs()
        f = jax.jit(self.f, static_argnames=('g', *static.keys()))
        return f(x, self.g, **(dynamic|static))
    
    @staticmethod
    def f(x: float, g: Callable, **kwargs) -> float:
        # Jitted parent method
        return kwargs['dp']*g(x, **kwargs) if kwargs['sp'] else g(x, **kwargs)
    
    @staticmethod
    def g(x):
        # To be implemented in child classes
        raise NotImplementedError
    
@dataclass
class Child(Parent):

    dc: float # dynamic child parameter
    sc: bool # static child parameter

    def kwargs(self):
        dynamic, static = super().kwargs()
        return (dynamic|dict(dc=self.dc), static|dict(sc=self.sc))

    @staticmethod
    def g(x, **kwargs):
        return kwargs['dc']*x if kwargs['sc'] else x

c = Child(2.0, False, 3.0, False)
%time c(1.0) # Compile
%time c(1.0) # No compile
print('\nChanging static args...')
c.sp = True  # Static arg change
%time c(1.0) # Compile
%time c(1.0) # No compile2
print('\nChanging static args...')
c.sc = True  # Static arg change
%time c(1.0) # Compile
%time c(1.0) # No compile
print('.\nChanging dynamic args...')
c.dp = 3.0   # Dynamic arg change
%time c(1.0) # No compile!
%time c(1.0) # No compile!
print('\nChanging dynamic args...')
c.dc = 4.0   # Dynamic arg change
%time c(1.0) # No compile!
%time c(1.0) # No compile!

I don't really like using a dictionary to handle this, but it seems to be a solution. What do you think? Does this already exist?

[jakevdp]

Yeah, that seems reasonable, though I'd probably prefer to wrap g in jit directly where it's defined

[cisprague]

Yeah, I would like to do the JITing where the method is defined, but it seems that the partitioning of the dynamic and static arguments is not accessible outside of the method definitions. Maybe we could use the data class field metadata as a workaround. It seems like something similar has been done in the dataclass-type decorators of Flax, Chex, and TJax. But I'm not sure if they would allow changes to attributes in the way I want.