return a ConstantMulLT instead
@@ -159,9 +159,7 @@ class ConstantMulLazyTensor(LazyTensor):
@cached(name="root_decomposition")
def root_decomposition(self, method: Optional[str] = None):
if (self._constant > 0).sum() == self._constant.numel():
-
base_root = self.base_lazy_tensor.root_decomposition(method=method).root
- constant_root = self.expanded_constant ** 0.5
- return RootLazyTensor(constant_root * base_root)
+ return RootLazyTensor(ConstantMulLazyTensor(base_root, self._constant ** 0.5))
return super().root_decomposition(method=method)
