argmin, argmax, argwhere

[lukasgebhard]

Hi!

No wonder I like einops: Without it, I'm always lost whenever I use a NumPy function with an axis argument... which axis was the one I need? 🤔

Among the NumPy functions that have an axis parameter are argmin, argmax, and argwhere. E.g.:

import numpy as np

array = np.array([
    [1,1,1],
    [0,2,0]
], dtype=float)
result = np.argmax(array, axis=0)
print(result)  # [0, 1, 0]

Is there a way to do the same thing with einops (and thus avoid the axis argument)? The following doesn't work:

import einops

einops.reduce(array, "x y -> y", np.argmax)  # TypeError: 'tuple' object cannot be interpreted as an integer

[arogozhnikov]

Hi @lukasgebhard

https://github.com/arogozhnikov/eindex aims to cover argmin/argmax/topk (and some other) operations.

It relies on array api, use most function can be applied to any array-api compatible lib (e.g. numpy 2.0 or new jax, but not torch).

[lukasgebhard]

Thanks for the pointer!

It would be nice if I could use the einops syntax:

import eindex.numpy as ex

result = ex.argmax(array, "x y -> y")  # eindex.EindexError: Composition axis should go first in indexer, like '[h, w] i j k', not 'y'

As this syntax is not supported, I'd have to write:

result = ex.argmax(array, "x y -> [x] y").squeeze(axis=0)
print(result)  # [0 1 0]

so the axis argument (which I hoped to avoid) would creep back in 😅

[arogozhnikov]

you're supposed to write it like:

[x] = ex.argmax(array, "x y -> [x] y")
[x, z] = ex.argmax(array, "x y z -> [x, z] y") # multidim argmax