Back out "Removing custom BlockDiagLazyTensor logic when using Standardize" (#1383)

Summary:
Pull Request resolved: #1383

With this change the `Multi_objective_multi_fidelity_BO` tutorial is taking much much longer to run (4s vs 50s). Reverting this to avoid unexpected regressions in other places.

Reviewed By: Balandat, SebastianAment

Differential Revision: D39361738

fbshipit-source-id: 831df5033c94e4d562313d9e932b0febc18235cb

Author: saitcakmak

botorch/models/transforms/outcome.py

@@ -33,7 +33,11 @@
 )
 from botorch.posteriors import GPyTorchPosterior, Posterior, TransformedPosterior
 from botorch.utils.transforms import normalize_indices
-from linear_operator.operators import CholLinearOperator, DiagLinearOperator
+from linear_operator.operators import (
+    BlockDiagLinearOperator,
+    CholLinearOperator,
+    DiagLinearOperator,
+)
 from torch import Tensor
 from torch.nn import Module, ModuleDict
 
@@ -382,8 +386,19 @@ def untransform_posterior(self, posterior: Posterior) -> Posterior:
         else:
             lcv = mvn.lazy_covariance_matrix
             scale_fac = scale_fac.expand(lcv.shape[:-1])
-            scale_mat = DiagLinearOperator(scale_fac)
-            covar_tf = scale_mat @ lcv @ scale_mat
+            # 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
 
         kwargs = {"interleaved": mvn._interleaved} if posterior._is_mt else {}
         mvn_tf = mvn.__class__(mean=mean_tf, covariance_matrix=covar_tf, **kwargs)