Mypy error when indexing into array with ellipsis and unpacking

Question

I get a mypy error when I index into an array using both Ellipsis and tuple unpacking.

I have a line of code that looks like this:

new_mat = mat[..., *np.ix_(inds, inds)]

which gives rise to the following mypy error

error: Argument 1 to <tuple> has incompatible type "EllipsisType"; expected "ndarray[tuple[int, ...], dtype[integer[Any] | numpy.bool[builtins.bool]]]" [arg-type]

I have similar lines of code several times in close proximity, with different matrices and sets of indices. I would prefer not to sprinkle mypy ignore messages everywhere, or otherwise increase the verbosity of my code. The code is also called with jax arrays (in which case np is really jax.numpy), including by jitted functions.

Is there an idiomatic way to resolve the mypy error?

• STerliakov: A wild guess: does running mypy with --enable-incomplete-feature=PreciseTupleTypes enabled help? If yes, just do so, if no - please report as a bug at github.com/python/mypy along with your mypy, python and numpy versions. Also please try with --enable-incomplete-feature=PreciseTupleTypes on master to avoid reporting something that might have been recently fixed.
• aklmn: @STerliakov, that worked! Thank you very much.

Answer 1

@STerliakov's comment resolved the issue. Running mypy with --enable-incomplete-feature=PreciseTupleTypes resolves the issue by telling mypy to infer more precise tuple types.