[Feature Request] __rmatmul__ for LazyTensors

[j-wilson]

🚀 Feature Request

Implement __rtmatmul__ for LazyTensors.

Motivation

LazyTensors must be converted to vanilla Tensors when used by functions that perform matrix multiplies between Tensors (left) and LazyTensors (right). For example, if A is a lazy and K is not then expressions such as A @ K @ A.T will fail.

Describe the solution you'd like
I do not know what challenges present themselves here. This may be as simple as adding something like

def __rmatmul__(self, other: Tensor) -> Tensor:
    if other.ndim == 1:
        return self.transpose(-1, -2).matmul(other)
    return self.transpose(-1, -2).matmul(other.transpose(-1, -2)).transpose(-1, -2)

Are you willing to open a pull request?
Potentially, yes. It would be good to determine the scope of the PR first however.

[jacobrgardner]

I think an implementation that basically does the above (although of course also transpose other) would be a perfectly reasonable first pass at this.

The only thing to note is that lazy tensors that represent symmetric positive definite matrices (which is the dominant use case in gpytorch of course) don't even need to have _transpose_nonbatch implemented, since obviously self._transpose_nonbatch(-2, -1) would just be the same LT as self.

That said, I just looked through all of the LTs again, and basically every one I see either implements _transpose_nonbatch or is already actually explicitly assuming the matrix is symmetric, e.g.,

gpytorch/gpytorch/lazy/toeplitz_lazy_tensor.py

Lines 52 to 53 in 4df3048

	def _transpose_nonbatch(self):
	return ToeplitzLazyTensor(self.column)

so the quick implementation may be all we need here.

[j-wilson]

Good catch regarding transposing other, I've updated the earlier __rmatmul__ method accordingly. From what you've said, it sounds like this simple solution may suffice for the time being. If so, I'll make a quick PR. Please let me know if you would like any test cases added over and beyond something like the following:

def _test_rmatmul(self, lhs):
    lazy_tensor = self.create_lazy_tensor().detach().requires_grad_(True)
    lazy_tensor_copy = lazy_tensor.clone().detach().requires_grad_(True)
    evaluated = self.evaluate_lazy_tensor(lazy_tensor_copy)

    res = lhs @ lazy_tensor
    actual = lhs @ evaluated
    self.assertAllClose(res, actual)

    grad = torch.randn_like(res)
    res.backward(gradient=grad)
    actual.backward(gradient=grad)
    for arg, arg_copy in zip(lazy_tensor.representation(), lazy_tensor_copy.representation()):
        if arg_copy.requires_grad and arg_copy.is_leaf and arg_copy.grad is not None:
            self.assertAllClose(arg.grad, arg_copy.grad, **self.tolerances["matmul"])

[jacobrgardner]

@j-wilson Yeah, both the implementation and the basic test looks pretty good to me 👍