`jax.numpy.sum` and `jax.numpy.mean` inconsistency on python lists

This works:

import jax.numpy as jnp

x = [jnp.ones((), dtype=jnp.float32), jnp.ones((), dtype=jnp.float32)]
jnp.sum(x)

and gives me DeviceArray(2., dtype=float32). But this doesn’t:

import jax.numpy as jnp

x = [jnp.ones((), dtype=jnp.float32), jnp.ones((), dtype=jnp.float32)]
jnp.mean(x)

and I got TypeError: data type not understood.

In numpy both would work:

import numpy as np

x = [np.ones((), dtype=np.float32), np.ones((), dtype=np.float32)]
np.mean(x)  # 1.0
np.sum(x)  # 2.0

Is this intended?

[hawkinsp]

Both sum and mean now raise errors at head!

[jakevdp]

From discussions with team members, I think the preferred appoach here would be to not accept lists as inputs to JAX functions (in particular, converting lists to arrays may lead to strange interactions with how we handle pytrees internally).

So to make the API more consistent, we should ensure that sum() and other routines raise an appropriate TypeError when passed a Python list.