Inconsistent type hints on _matmul() and _t_matmul()

[Turakar]

These are the type hints on these two methods:

def _matmul(
    self: Float[LinearOperator, "*batch M N"],
    rhs: Union[Float[torch.Tensor, "*batch2 N C"], Float[torch.Tensor, "*batch2 N"]],
) -> Union[Float[torch.Tensor, "... M C"], Float[torch.Tensor, "... M"]]:
    ...

def _t_matmul(
    self: Float[LinearOperator, "*batch M N"],
    rhs: Union[Float[Tensor, "*batch2 M P"], Float[LinearOperator, "*batch2 M P"]],
) -> Union[Float[LinearOperator, "... N P"], Float[Tensor, "... N P"]]:
    ...

As you can see, there are two inconsistencies:

1. _matmul() accepts both tensors with and without a trailing dimension, while _t_matmul() does not.
2. _t_matmul() accepts both a Tensor or a linear operator, while _matmul() does not.

Where are these inconsistencies coming from?

[gpleiss]

Hmm good catch. I think the correct type hint should be:

def _matmul(
    self: Float[LinearOperator, "*batch M N"],
    rhs: Union[Float[LinearOperator, "*batch2 M C"], Float[torch.Tensor, "*batch2 N C"], Float[torch.Tensor, "N"]],
) -> Union[Float[LinearOperator, "... N P"], Float[torch.Tensor, "... M C"], Float[torch.Tensor, "M"]]:

and analogously for _t_matmul.
(We shouldn't allow batch matmuls/t_matmuls for any batched linear operators: the matrix-vector case is only for non-batch matrices multiplied by vectors, following the api of torch's matmul.)

I'll fix this in a PR.

[Turakar]

I think the LinearOperator for rhs should have dimensions "*batch2 N C". Apart from that, this seems fine to me.