Passing 4 arguments to a jit function with 4 parameters raises TypeError: jit() takes 1 positional argument but 5 positional arguments [closed]

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?