HOWTO: differentiable functions that depend on static arguments?

[smao-astro]

Hi,

I am new to JAX and I am looking for suggestions that which is the best way to generate pure functions that depend on static arguments.

For example (over-simplified), I want to get a function f(x), that do math calculation: a*x**2+b*x+c, and a, b, c here are constants and does not change in one execution (depend on the command line input to the program); and I also want to have easy-to-write-and-read jacobians of f(x), such as

jax.vmap(jax.jacfwd(f))(x)

without considering the a, b, c.

I have got two way to do so, but I am wondering: 1. what is the difference of these two way. 2. which one is better. Please give me some help!

The first way:

def get_f(a):
    @jax.jit
    def f(x):
        return x + a
    return f

f = get_f(1)

The second way:

def g(x, a):
    return x + a

g = jax.jit(jax.tree_util.Partial(g, a=1))

Thanks!

[cgarciae]

Partials are OK. You can also create a Module-like structure:

@dataclass
class G:
    a: float
    def __call__(self, x):
        return self.a + x

g = jax.jit(G(a=1))