einx.{max,min} are broken with pytorch

[falkaer]

Hi,

Using einx.{max,min} with pytorch tensors return a namedtuple of type torch.return_types.{min,max} due to directly calling torch.{min,max} with a dim argument. This is an unexpected result, and also breaks some otherwise valid reduce/rearrange combinations such as:

import einx
import torch

x = torch.randn(5, 3)
print(einx.max("a [b] -> a 1", x))

Traceback (most recent call last):
  File "/home/falkaer/projects/latent-features/einx_mwe.py", line 5, in <module>
    print(einx.max("a [b] -> a 1", x))
          ^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/falkaer/projects/latent-features/.devenv/state/venv/lib/python3.11/site-packages/einx/traceback_util.py", line 71, in func_with_reraise
    raise e.with_traceback(tb) from None
  File "<string>", line 4, in op0
TypeError: reshape(): argument 'input' (position 1) must be Tensor, not torch.return_types.max

[fferflo]

Hi, thanks for bringing this up. einx.{min|max} should indeed use torch.{amin|amax} rather than torch.{min|max}. This is fixed in version 0.4.0 which will release in the next weeks.

[fferflo]

v0.4.0 is now released, and your code

import einx
import torch

x = torch.randn(5, 3)
print(einx.max("a [b] -> a 1", x))

runs without error. The compiled operation now calls torch.amax rather than torch.max:

>>> print(einx.max("a [b] -> a 1", x, graph=True))
import torch
def op(a):
    a = torch.amax(a, dim=1)
    a = torch.reshape(a, (5, 1))
    return a

Closing this issue, feel free to reopen if anything isn't working as expected.