[BUG] einops.repeat returns value with type Never

[adamjstewart]

Describe the bug

This is a weird one. When I pass a PyTorch Tensor into einops.repeat, mypy 1.7+ with strict mode enabled tries to tell me that the return type is typing.Never.

Reproduction steps

Create the following test.py file:

import einops
import torch

x = torch.rand(2, 2)
reveal_type(x)
x = einops.repeat(x, "h w -> h w c", c=3)
reveal_type(x)

Then run:

> mypy --strict test.py
test.py:5: note: Revealed type is "torch._tensor.Tensor"
test.py:7: note: Revealed type is "Never"

Expected behavior

I would expect the output to have the same type as the input, a PyTorch Tensor. This works fine for:

• numpy arrays
• non-strict mypy
• older mypy

I haven't checked any other einops functions.

Your platform Version of einops, python and DL package that you used

• einops 0.7.0
• PyTorch 2.1.2
• mypy 1.8.0
• Python 3.11.6

I'm trying to figure out if this is a bug in einops, PyTorch, or mypy. Any help you can give me would be appreciated. I've noticed several PyTorch-specific issues with mypy 1.7+, so I doubt this is just an einops issue.

[arogozhnikov]

Hi Adam,

einops.repeat typing logic is literally in this line: https://github.com/arogozhnikov/einops/blob/d0c7feef31eed7adceff07175300ebfc8bdee2b2/einops/einops.py#L459

and it wasn't changed for a long time. Also it is agnostic to framework, so I assume that's something with mypy.

[adamjstewart]

Can confirm that this issue was fixed in PyTorch 2.2!

[arogozhnikov]

that's great, thanks for leaving this note