Typo in the permissive_sum "gotcha"?

[lrizzello]

Hello,

I started looking into jax today, and while reading the common gotchas notebook I ran into something that confused me.

Towards the 14th line, non array inputs are discussed. We are given a bad example of giving a list to a jax function here

def permissive_sum(x):
  return jnp.sum(jnp.array(x))

x = list(range(10))
permissive_sum(x)

We are then told that a better way of doing this would be
jnp.sum(jnp.array(x))
But, except for the fact that this line of code is standalone and not in a function, it is the same line.

Is this a typo or is there something I missed? I assume the first example was supposed to be jnp.sum(x).

[soraros]

I think the point they are trying to make is that you shouldn't pass in lists and convert them to arrays inside of a function.
You should read jnp.sum(jnp.array(x)) as passing an array jnp.array(x) to the function jnp.sum(for there is no need to wrap it in a lambda).

[lrizzello]

I just reread the notebook, and saw the line make_jaxpr(permissive_sum)(x) (not sure why I didn't pay attention to that line the first time around).
If I'm understanding this correctly, this is the issue. It turns the entire permissive_sum function into a jax expression and, as such, this will result in a jax expression that will accept a list as input, which is inefficient. Is that correct?

[jakevdp]

If x is a list, jnp.sum(x) will raise an error. The permissive_sum code you refer to is an example used in describing the rationale for that API choice.

The workaround is to convert lists to arrays before passing them to JAX functions, i.e. jnp.sum(jnp.array(x)) rather than jnp.sum(x).

[jakevdp]

Maybe it would be clearer if we used something like this?

def permissive_sum(x):
  return jnp.asarray(x).sum()