JIT often recompiles bound methods

I’m going close this issue. I just want to save a record of it in case I come up with a reasonable solution one day.

If a method is always jitted, then Jax compiles it once:

from jax import jit
from functools import partial

class X:
  @partial(jit, static_argnums=(0,))
  def f(self):
    print("Recompiling")

x = X()
for i in range(3):
    x.f()

However, if the caller wants to jit the method, then Jax compiles it every time:

from jax import jit

class X:
  def f(self):
    print("Recompiling")

x = X()
for i in range(3):
    jit(x.f)()

It would be nice if jit detected the case when the function f is a bound method (use inspect.ismethod(f)), and maybe emit a warning.

If you wanted to fix it automatically, you could use replace the function with function.__func__, and add a positional argument function.__self__. The issue with this is whether that argument should be static or not is impossible to know. That’s why the warning makes sense: it forces the caller to explicitly choose whether the argument is static or dynamic.

[patrick-kidger]

team

Just me! 😄

[jakevdp]

FYI, there is now a section in the FAQ addressing jit compilation of methods: https://jax.readthedocs.io/en/latest/faq.html#how-to-use-jit-with-methods