Double jtu.Partial results in cache miss

[patrick-kidger]

Description

import jax
import jax.tree_util as jtu

def f(x): pass
g = jtu.Partial(f)
h1 = jtu.Partial(g)
h2 = jtu.Partial(g)

def check(x):
    print("tracing!")

check = jax.jit(check, static_argnums=0)
check(h1)  # tracing!
check(h2)  # tracing!

Both h1 and h2 have the same PyTree structure so expected behaviour is for the second call not to print out anything.

What jax/jaxlib version are you using?

0.3.23

[cgarciae]

Hey @patrick-kidger, because you are passing these as static and Partial seems to implement hash by object id (or something similar):

hash(h1) != hash(h2)

I'd say this behaviour is "expected".

[patrick-kidger]

I don't think so. The single-level partial works just fine, despite the fact that each Partial object will have different object ids.

I think it is the special handling here that's misbehaving:

https://github.com/google/jax/blob/96f6c1c9d414e6ebc54ff7f08115a9a9a6d6a8f8/jax/_src/tree_util.py#L366

[KeAWang]

I agree with @cgarciae actually. This seems like not a bug? I tried your single-level partial case and it seems to behave as expected:

import jax
import jax.tree_util as jtu

def f(x):
    return 

g1 = jtu.Partial(f)
g2 = jtu.Partial(f)

assert hash(g1) != hash(g2)

def check(x):
    return

with jax.log_compiles(True):
    check1 = jax.jit(check, static_argnums=0)
    check1(g1) # compiles since it's the first call
    check1(g2)  # compiles since hash(g1) != hash(g2)

with jax.log_compiles(True):
    check2 = jax.jit(check)
    check2(g1)  # tracing!
    check2(g2)  # no tracing as expected since now g1 and g2 are treated as dynamic

[patrick-kidger]

Right; I've since recalled that the function-valued argument to jtu.Partial is assumed to be static. Thus Partial(Partial(...), ...) is actually a user error, since the inner Partial(...) isn't really a static thing.

I think there should simply probably be a check to guard against this error, e.g. raise an error if len(jtu.tree_leaves(func)) > 1.

And FWIW in my own code, for this reason I no longer use jtu.Partial, and use my own implementation instead:

class Partial(eqx.Module):
    func: Callable
    args: tuple
    kwargs: dict

    def __init__(self, func, *args, **kwargs):
        self.func = func
        self.args = args
        self.kwargs = kwargs

    def __call__(self, *args, **kwargs):
        return self.func(*self.args, *args, **kwargs, **self.kwargs)

[KeAWang]

That's a good check. And maybe add some documentation to its assumptions too...