Thanks for making this library; I wanted to do something similar a few months ago but other things got into the way; awesome to see that the JAX ecosystem is maturing this fast!

One design question I had was wether to go with a struct-of-arrays or array-of-structs memory layout. Unless ive misread, your design is the latter; that is, if I vmap over a multivector, the components of the multivector will form the last axis; which in JAX is the contiguous stride==1 axis. If youd think about how this would get vectorized on a GPU, that wouldnt be ideal; if every thread in a block gets to work on a single element of the vmapped axis, which is the most straightforward parrelelization, now the threads in this warp are not performing contiguous memory accesses. Hence the structs-of-arrays are generally preferred on the GPU. Also if you dig into the deepmind/alphafold repo, you will see that they also use a struct-of-array layout for their vector types and the like.

Now this is all terrible premature optimization as far as the actual goals im trying to achieve; but I guess im trying to form a bit of a deeper understanding of JAX and TPUs on a low level. So with that in mind; was this a deliberate choice, or something that you have given any thought to?

• Previously was using a loop and add at index which gets unrolled, now using segment_sum to sum same output indices
• 10x faster JIT on CPU, 6x faster JIT on GPU
• 100x slower runtime on CPU, 5x faster runtime on GPU

Should maybe add a flag for whether to use this one, very useful for large algebras where JIT takes very long because of analyzing the unrolled loop. Maybe make it the default on GPU too.

CPU results show segment_sum runtime very dependent on batch size

a_val, a_ind = jnp.array(jnp.ones([5, 10]), dtype=jnp.float32), tuple((i,) for i in range(5))
b_val, b_ind = jnp.array(jnp.ones([5, 10]), dtype=jnp.float32), tuple((i,i+1) for i in range(5))

new:
Wall time: 94 ms
10.8 µs ± 301 ns per loop (mean ± std. dev. of 7 runs, 100000 loops each)

old:
Wall time: 1.04 s
10.9 µs ± 152 ns per loop (mean ± std. dev. of 7 runs, 100000 loops each)

---
a_val, a_ind = jnp.array(jnp.ones([5, 100]), dtype=jnp.float32), tuple((i,) for i in range(5))
b_val, b_ind = jnp.array(jnp.ones([5, 100]), dtype=jnp.float32), tuple((i,i+1) for i in range(5))

new:
Wall time: 227 ms
48.6 µs ± 640 ns per loop (mean ± std. dev. of 7 runs, 10000 loops each)

old:
Wall time: 676 ms
11.6 µs ± 464 ns per loop (mean ± std. dev. of 7 runs, 100000 loops each)
---
a_val, a_ind = jnp.array(jnp.ones([10, 100]), dtype=jnp.float32), tuple((i,) for i in range(5))
b_val, b_ind = jnp.array(jnp.ones([10, 100]), dtype=jnp.float32), tuple((i,i+1) for i in range(5))

new:
Wall time: 261 ms
49.1 µs ± 1.49 µs per loop (mean ± std. dev. of 7 runs, 10000 loops each)

old:
Wall time: 687 ms
11.6 µs ± 196 ns per loop (mean ± std. dev. of 7 runs, 100000 loops each)
---
a_val, a_ind = jnp.array(jnp.ones([10, 1000]), dtype=jnp.float32), tuple((i,) for i in range(5))
b_val, b_ind = jnp.array(jnp.ones([10, 1000]), dtype=jnp.float32), tuple((i,i+1) for i in range(5))

new:
Wall time: 256 ms
1.19 ms ± 69.3 µs per loop (mean ± std. dev. of 7 runs, 1000 loops each)

old:
Wall time: 558 ms
16.9 µs ± 234 ns per loop (mean ± std. dev. of 7 runs, 100000 loops each)