Compared to list，the operation 'at' of jax array may take longer ？

[jing-alice]

I want to get better performance with vmap. This is the code.

for i in range(len(self.conn_list)):
        conn = self.conn_list[i]
        pre_idx = self.conn_pre_idx[i]
        post_idx = self.conn_post_idx[i]
        v_pre = var_list[pre_idx]['v']
        v_post = var_list[post_idx]['v']
        conn_I = conn(v_pre, v_post)
        _input[pre_idx] -= conn_I
        _input[post_idx] += conn_I

Also, the for has JIT.
And I have tried vmap in the follwing:

_input = jnp.zeros_like(_flatten_var)

       def conn_f(pre_idx, post_idx, _input):
        v_pre = _flatten_var[pre_idx]
        v_post = _flatten_var[post_idx]
        conn = 10. * (v_pre - v_post)
        _input = _input.at[pre_idx].add(-conn)
        _input = _input.at[post_idx].add(conn)   
        return _input
conn_input = jnp.squeeze(
  vmap(conn_f, in_axes=(0, 0, None))(jnp.array(self.conn_pre_idx), jnp.array(self.conn_post_idx), _input))
_input += jnp.sum(conn_input, axis=0)

In vmap, we must use jax array, but it's slower.
The first picture is for,

the second picture is vmap

.

so compared to list，the operation 'at' of jax array may take longer ？Is there a better way for jax to avoid this problem?

[jakevdp]

Hi, thanks for the question. There's a section of the FAQ that's relevant: FAQ: Is JAX Faster Than NumPy. Although this focuses on comparing JAX to NumPy, the discussion is relevant here as well.

In short, JAX has relatively expensive per-operation dispatch costs. This doesn't really matter in practice, because in typical use you pay that dispatch cost only once per (JIT-compiled) program.

By contrast, Python has very low per-operation dispatch cost. This is important because in Python you must pay this cost for every operation in your program: there is no built-in JIT compiler to get around that.

Given this, when you compare very cheap operations like a single indexing op, you're essentially just comparing the single-operation dispatch cost, for which we know JAX will be slower than raw Python. But that has very little bearing on how JAX will perform in real-world use-cases, which generally involve longer sequences of JIT-compiled operations.