[question] how to properly use lax.reduce on an multi-dimension array.

[linzzzzzz]

I want to perform reduce operation on an array in a slightly different way than jnp.argmax.

I have an array A with shape (2,5). I wanted to reduce A in the below way: for a pair of index (i, j), calculate z = A[0][i]*A[1][j]-A[0][j]*A[1][i], if z>0, then pick i otherwise pick j. Continue this computation until getting the final index.

What's the best way for me to perform this in JAX?

Below is my code on trying to implement this reduce operation, which threw error IndexError: Too many indices for array: 1 non-None/Ellipsis indices for dim 0..

A = np.random.randn(2, 5)

def my_compare(x0, x1):
    y0 = x0[0]*x1[1]
    y1 = x0[1]*x1[0]
    z = y0 > y1
    return jax.lax.select(z, x0, x1)

res = jax.lax.reduce(A, A[0][0], my_compare, [0,1])

[jakevdp]

You cannot use lax.reduce for the operation you have in mind. lax.reduce is a low-level operator that has pretty stringent requirements for the function you pass to it: in particular, it must form a monoid with the initial value, and I don't think that's true of the function you're using (there's a brief discussion in the docs: https://jax.readthedocs.io/en/latest/_autosummary/jax.lax.reduce.html)

I don't immediately see any symmetries in this problem that would lend themselves to an efficient vectorized operation, so probably the best option here would be to implement it in terms of lax.scan or lax.fori_loop; here's an example that I think implements the logic you have in mind:

A = jnp.asarray(A)
result = jax.lax.fori_loop(1, A.shape[1], lambda i, val: my_compare(val, A[:, i]), A[:, 0])

[linzzzzzz]

Thanks @jakevdp for answering.

I read more on the link you posted as well as the background on monoid. Do you think Variadic Reduce of lax.reduce might help? For example, I can split array A with shape (2,5) into two separate array each with shape (,5).

[jakevdp]

Unless you’re reducing over a monoid, you won’t be able to use lax.reduce no matter how you slice it.