JAX extremely slow under some circumstances (MWE included)

[JoaoAparicio]

JAX takes about 30x longer than numpy in converting lists to its own arrays.

I've attached a file here. The code would be:

import json

import numpy as np
import jax.numpy as jnp

def from_json(state):
    if isinstance(state, list):
        return np.array(state)
    elif isinstance(state, dict):
        out = {}
        for k in state:
            out[k] = from_json(state[k])
        return out
    else:
        return state 

def jax_from_json(state):
    if isinstance(state, list):
        return jnp.array(state)
    elif isinstance(state, dict):
        out = {}
        for k in state:
            out[k] = jax_from_json(state[k])
        return out
    else:
        return state


with open(path, "r") as f:
    d = json.load(f)

%timeit from_json(d)  # aprox 90us per

%timeit jax_from_json(d)  # aprox 3ms per

I guess that in here this would qualify as a "micro-benchmark". However I think this minimizes the problem. If I'm calling this in a loop because that's how I get the data and it becomes significant compared to the rest of the problem, it's no longer micro.

I understand that JAX does a lot of difficulty things that numpy can't (JAX is amazing, thank you), but at the same time if it's going to be mentioned as "can often be used as drop-in replacement" to numpy (e.g. here), it shouldn't be 30x slower at any point...

Alternatively it's also possible I'm doing something wrong, in which case help? :-)

jax==0.4.37
jaxlib==0.4.36
jaxtyping==0.2.36