Please allow @jax.jit on staticmethods

The following code does not work,

class Foo:
  @jax.jit
  @staticmethod
  def foo(x):
    return x * 2

with an error TypeError: Expected a callable value, got <staticmethod object at 0x7fe007e767f0>. This is because a staticmethod is NOT a callable, until Python 3.10. Of course, staticmethod is callable – something we can call.

In #1251 one possible workaround is suggested:

class Foo:
  @functools.partial(jax.jit, static_argnums=(0,))
  def foo(self, x):
    # this function should not make use of `self` to be pure-functional
    return x * 2

which is a bit ugly, but basically does the same thing as long as self is not used.

I don’t see any reason staticmethod cannot be jit-ed as other non-method functions. I think this is just a matter of extending _check_callable(fun) to support staticmethods (which will happen automatically in Python 3.10); any pitfalls I would’ve missed?

[jakevdp]

I think the cleanest approach is probably to use @staticmethod @jit rather than @jit @staticmethod.

If what you say about Python 3.10 is correct, though, the tests referred to above will (I think?) fail when run with Py3.10. If there are no fundamental reasons why staticmethod objects cannot be jitted, I would be fine with a PR that removes those tests and updates _is_callable to recognize static methods.

[wookayin]

Thank you @jblespiau for the explanation.

It’s always possible to have a plain, top-level function doing the same thing. So usually, we never have static methods, and thus, never want to jax.jit them.

That would probably be how Google usually writes a JAX code, but I respectfully disagree staticmethods should be prohibited, because some community codes that do not necessarily follow the Google convention strictly might still want jit-ed staticmethods. Personally I don’t like having a top- or module-level plain function because I have to place them outside the class, usually very far from the methods that implement a very relevant logic.

That said, as an alternative one could define a (nested) plain function inside __init__ and assign them as an attribute, but I feel using staticmethods would give a bit more flexibility. I can prepare a PR to support this feature.