Removing custom BlockDiagLazyTensor logic when using Standardize (#1376)

Summary:
Pull Request resolved: #1376

Removing custom logic for `BlockDiagLazyTensor` when using `Standardize`, in particular in the `untransform_posterior` function. The special cases are now handled by [a meta-constructor of `BlockDiagonalLinearOperator`](cornellius-gp/linear_operator#4) and the `mul` method.

Reviewed By: Balandat

Differential Revision: D38242402

fbshipit-source-id: 11461db1be1d2c0581571693406e2d54aa7110d5

Author: SebastianAment

botorch/models/transforms/outcome.py

@@ -33,11 +33,7 @@
 )
 from botorch.posteriors import GPyTorchPosterior, Posterior, TransformedPosterior
 from botorch.utils.transforms import normalize_indices
-from linear_operator.operators import (
-    BlockDiagLinearOperator,
-    CholLinearOperator,
-    DiagLinearOperator,
-)
+from linear_operator.operators import CholLinearOperator, DiagLinearOperator
 from torch import Tensor
 from torch.nn import Module, ModuleDict
 
@@ -386,19 +382,8 @@ def untransform_posterior(self, posterior: Posterior) -> Posterior:
         else:
             lcv = mvn.lazy_covariance_matrix
             scale_fac = scale_fac.expand(lcv.shape[:-1])
-            # TODO: Remove the custom logic with next GPyTorch release (T126095032).
-            if isinstance(lcv, BlockDiagLinearOperator):
-                # Keep the block diag structure of lcv.
-                base_lcv = lcv.base_linear_op
-                scale_mat = DiagLinearOperator(
-                    scale_fac.view(*scale_fac.shape[:-1], lcv.num_blocks, -1)
-                )
-                base_lcv_tf = scale_mat @ base_lcv @ scale_mat
-                covar_tf = BlockDiagLinearOperator(base_linear_op=base_lcv_tf)
-            else:
-                # allow batch-evaluation of the model
-                scale_mat = DiagLinearOperator(scale_fac)
-                covar_tf = scale_mat @ lcv @ scale_mat
+            scale_mat = DiagLinearOperator(scale_fac)
+            covar_tf = scale_mat @ lcv @ scale_mat
 
         kwargs = {"interleaved": mvn._interleaved} if posterior._is_mt else {}
         mvn_tf = mvn.__class__(mean=mean_tf, covariance_matrix=covar_tf, **kwargs)