minor typing fixes in einops.py

[jorenham]

Note that TYPE_CHECKING is always False at runtime, so numpy will never actually be imported. This if TYPE_CHECKING just ensures that we can use the np.ndarray as static type annotation here, and avoid having to unnecessarily useAny (it's maximally type-unsafe).

I like the new TensorLike protocol. Before, it required that the __getitem__ parameter (arg) was actually named "arg". By changing it to a positional-only parameter, it will now allow any name. This doesn't mean that it's now required for types to also have a pos-only parameter for __getitem__, because of Barbara Liskov's substitution principle.

This also introduces a private _Axis type alias, which is used in several new annotations that resolve at least 4 mypy errors.

[arogozhnikov]

I'm working a patch that addresses these + a bit more, hopefully will get time over the weekend to finish

[jorenham]

Ah even better. Feel free to ping me if you want me to give it a quick scan or somehting once you get there.