[Python] Passing 4 arguments to a jit function with 4 parameters raises TypeError: jit() takes...

I am currently working on GPU optimized simulations. For that, I have a function where the head looks like this:

@partial(partial, jax.jit, static_argnums=(2,4))
def init_particles_grid(n_particles: int, n_side: int, L: float, A: float):


Once I call it:

n_side = 4
N = n_side * n_side
L = 1.0
A = 1.0 / (4.0 * jnp.pi**2)
q, x0, v0 = init_particles_grid(N, n_side, L, A)


I get

TypeError: jit() takes 1 positional argument but 5 positional arguments (and 1 keyword-only argument) were given
​

What am I doing wrong?

Continue reading...