`jnp.array` is much slower than `np.array` when called on a nested list

Consider the following example (I am using JAX version 0.3.10):

>>> import jax.numpy as jnp
>>> import numpy as np
>>> x = [list(range(1000)) for i in range(1000)]    # A simple nested list
>>> %timeit np.array(x)
79.4 ms ± 6.85 ms per loop (mean ± std. dev. of 7 runs, 10 loops each
>>> %timeit jnp.array(x)
3.22 s ± 174 ms per loop (mean ± std. dev. of 7 runs, 1 loop each)

NumPy is orders of magnitude faster than JAX here. A similar problem for non-nested lists (#2919) was addressed in PR #3350.

[jakevdp]

Sure, one important reason we can’t just call np.array when the argument is a list is that JAX transforms inject tracers in place of values, so np.array will result in an error:

from jax import jit

@jit
def f(x):
  return jnp.array([x, 0, 0])
print(f(x))
# [0. 0. 0.]

@jit
def g(x):
  return np.array([x, 0, 0])
g(x)
# ---------------------------------------------------------------------------
# TracerArrayConversionError: The numpy.ndarray conversion method __array__() was called on the JAX Tracer object Traced<ShapedArray(float16[])>with<DynamicJaxprTrace(level=0/1)>
# While tracing the function g at <ipython-input-13-53206b980741>:9 for jit, this concrete value was not available in Python because it depends on the value of the argument 'x'.
# See https://jax.readthedocs.io/en/latest/errors.html#jax.errors.TracerArrayConversionError

Another more subtle example is when you have a list with a mix of dtypes where type promotion behavior differs between numpy and jax:

import jax.numpy as jnp
import numpy as np

x = jnp.float16(0)

print(np.array([x, 0, 0]).dtype)
# float64

print(jnp.array([x, 0, 0]).dtype)
# float16

JAX treats unspecified python scalars as “weakly-typed”, in that they conform to the types of other “strongly-typed” values they interact with. Numpy treats Python scalars as 64-bit, and tends to promote things to 64-bit indiscriminately. This is problematic when working on accelerators like GPU or TPU, where it is often preferable to work in lower-precision for performance reasons, but users still want the convenience of writing things like 0 and 1 without explicit dtype declaration.

[jakevdp]

I’ll note also that even the previous “fast path” that you reference in that PR would slow down by orders of magnitude once jit-compiled, and in many cases would return different results between jit and non-jit execution, which both are behaviors we try to avoid.