Custom root_inv_decomposition for DiagLazyTensor
@@ -155,6 +155,9 @@ class DiagLazyTensor(LazyTensor):
def root_decomposition(self):
return RootLazyTensor(DiagLazyTensor(self._diag.sqrt()))
+ def root_inv_decomposition(self):
+ return RootLazyTensor(DiagLazyTensor((1. / self._diag).sqrt()))
+
def zero_mean_mvn_samples(self, num_samples):
if self.ndimension() == 3:
base_samples = torch.randn(
