Thanks for sharing!
Yes, JAX could be much faster with JIT on a CPU. But even with JIT compiler, when the dimension is so large that we can ignore the CPU overhead, the manual mode (torch.man) is faster than JAX.
I have updated my code 4 to measure the time with jax.jit and jax.vmap.
jac_fun = jax.jit(jax.jacrev(model))
Jx_fn = jax.jit(jax.vmap(jac_fun, in_axes=(0,)))

On a GPU (NVIDIA Tesla V100),
mode: torch.auto
bs: 32
n_inputs: 1000
hidden_ndim: 1000
n_outputs: 1000
n_layers: 3
-------------
loop: 100
device: cuda
torch auto rev: 2.591s

mode: torch.man
bs: 32
n_inputs: 1000
hidden_ndim: 1000
n_outputs: 1000
n_layers: 3
-------------
loop: 100
device: cuda
torch manual rev: 0.454s

mode: jax
bs: 32
n_inputs: 1000
hidden_ndim: 1000
n_outputs: 1000
n_layers: 3
-------------
loop: 100
device: [GpuDevice(id=0)]
jit(jax.jacrev): 1.024s

This is not about PyTorch vs JAX, but the FLOPS we need to evaluate the Jacobian.
It is the nature of the auto-grad to evaluate the vector-Jacobian product (vjp) or the Jacobian-vector product (jvp), so you need extra computation compared to the manual-mode.
for the reverse-mode, auto-grad requires extra (M_L)(M_1)(M_0) FLOPS.
M_L: output dimension of NN
M_0: input dimension of NN (the first layer)
M_1: output dimension of the first layer