Performant Scaling of BlockDiagLinearOperator by DiagLinearOperator

[SebastianAment]

The primary goal of this PR is to enable the efficient scaling of BlockDiagLinearOperators by DiagLinearOperators. This will allow us to remove a special case in BoTorch's outcome transform.

In order to achieve this, this PR modifies and adds special cases to DiagLinearOperator and BlockDiagLinearOperator's matmul.

I tested the notebook that exhibited the regression by patching in the function definitions here and ended up with a 16 second runtime - as opposed to 30+ minutes before.

[SebastianAment]

This line was the source of the performance regressions. The reasoning behind it appears to be because MulLinearOperator always computes a root decomposition, which is both inefficient and introduces dead code in its implementation (see below). I am sidestepping this by replacing the * with secondary @ operators in the new special cases of the DiagLinearOperator and BlockDiagLinearOperator matmul methods, leading to MatmulLinearOperators instead.

However, this does not get rid of the more general issue. To fix that, I propose two steps in a future PR:

1. Introducing logic in the constructor of MulLinearOperator that decides whether or not to build a root decomposition.
2. Even if a root decomposition seems beneficial, delaying its computation until the very last moment when it is needed in matmul and cache the result. This will give us ~0 overhead in the case where the linear operator represents a posterior covariance matrix that is constructed via a posterior call but only the posterior mean is needed, as was the case in the notebook that exhibited the regression.

[gpleiss]

I think this change should be fine, since the MulLinearOpeator constructor performs root decompositions on left_linear_op and right_linear_op:

linear_operator/linear_operator/operators/mul_linear_operator.py

Line 25 in caead4d

if not isinstance(left_linear_op, RootLinearOperator):

Even if a root decomposition seems beneficial, delaying its computation until the very last moment when it is needed in matmul and cache the result.

Agreed.

[Balandat]

Overall this lgtm but I'll let @gpleiss double check the broader implications of this change.

[gpleiss]

I think this change should be fine, since the MulLinearOpeator constructor performs root decompositions on left_linear_op and right_linear_op:

linear_operator/linear_operator/operators/mul_linear_operator.py

Line 25 in caead4d

if not isinstance(left_linear_op, RootLinearOperator):

Even if a root decomposition seems beneficial, delaying its computation until the very last moment when it is needed in matmul and cache the result.

Agreed.