Decorator factory pattern

I couldn’t find a pre-existing issue for this, but I figure it might be useful to add support for the decorator factory pattern for the most common transformations (grad, jit, vmap, pmap, etc).

Currently people wanting to decorate functions but also pass parameters to the transformation have to either do:

from functools import partial
from jax import jit

@partial(jit, param1=..., param2=...)
def foo(bar):
	...

Or

from jax import jit

def foo(bar):
	...

foo = jit(foo, param1=..., param2=...)

This seems needlessly verbose, but I am wondering if perhaps there is some rationale behind this that I might not be aware of. Being used to numba’s jit decorator, I would expect the decorator factory pattern to work.

I think it would be super neat if the following pattern was supported:

from jax import jit

@jit(param1=..., param2=...)
def foo(bar):
	...

[jakevdp]

#10074 adds some documentation of the @partial(jit, ...) pattern in a couple places.

[msegado]

Came here to suggest this and found @JeppeKlitgaard had beat me to it… don’t want to bikeshed though so please close this if you prefer to focus on other things now that the current behavior is documented!

That said, one argument in favor of supporting the decorator-or-factory pattern is that the Python standard library does so in several places:

• https://docs.python.org/3/library/dataclasses.html#dataclasses.dataclass (<= particularly prominent)
• https://docs.python.org/3/library/functools.html#functools.lru_cache
• https://docs.python.org/3/library/xmlrpc.server.html#xmlrpc.server.SimpleXMLRPCServer.register_function

I’d argue that the convenience is worth it given that the jax team itself makes use of both @dataclass and @dataclass(...) in their source code 🙂