[Python] JIT: partial or with static argnums? Non hashable input, but hashable partial

I am a bit lost on what exactly going on and what option to choose. Let's go trough an example:

import jax
from functools import partial
from typing import List

def dummy(a: int, b: List[str]):
return a + 1



As b argument is mutable, jitting with static argnames will be failed:

j_dummy = jax.jit(dummy, static_argnames=['b'])
j_dummy(2, ['kek'])
ValueError: Non-hashable static arguments are not supported


However, if we do partial: jp_dummy = jax.jit(partial(dummy, b=['kek'])), we aim the goal. Somehow, partial object is indeed has __hash__ method, so we can check it with hash(partial(dummy, b=['kek'])).

So, I am a bit lost here: how I should proceed in a bigger picture? Should I produce partial functions with whatever arguments and then jit them or should I try to maintain my arguments hashable? What are situations when one approach is better than other? Is there any drawbacks?

Continue reading...