Support batches of constants in MulLinearOperator._mul_constant

linear_operator/operators/mul_linear_operator.py

@@ -82,7 +82,7 @@ def _matmul(
     def _mul_constant(
         self: LinearOperator, other: float | torch.Tensor  # shape: (*batch, M, N)
     ) -> LinearOperator:  # shape: (*batch, M, N)
-        if other > 0:
+        if (other > 0).all():
             res = self.__class__(self.left_linear_op._mul_constant(other), self.right_linear_op)
         else:
             # Negative constants can screw up the root_decomposition

linear_operator/operators/triangular_linear_operator.py

@@ -112,7 +112,7 @@ def _matmul(
     def _mul_constant(
         self: LinearOperator, other: float | torch.Tensor  # shape: (*batch, M, N)
     ) -> LinearOperator:  # shape: (*batch, M, N)
-        return self.__class__(self._tensor * other.unsqueeze(-1), upper=self.upper)
+        return self.__class__(self._tensor * other.unsqueeze(-1).unsqueeze(-1), upper=self.upper)
 
     def _root_decomposition(
         self: LinearOperator,  # shape: (..., N, N)

test/operators/test_mul_linear_operator.py

@@ -1,10 +1,11 @@
 #!/usr/bin/env python3
 
+import itertools
 import unittest
 
 import torch
 
-from linear_operator.operators import RootLinearOperator
+from linear_operator.operators import DenseLinearOperator, MulLinearOperator, RootLinearOperator
 from linear_operator.test.linear_operator_test_case import LinearOperatorTestCase
 
 
@@ -13,6 +14,26 @@ def make_random_mat(size, rank, batch_shape=torch.Size(())):
     return res
 
 
+class BatchConstantMulMixin:
+    def test_constant_mul_batch(self):
+        batch_shape = self.create_linear_op().batch_shape
+        mats = [torch.randn(*batch_shape, 6, 6) for _ in range(2)]
+        mul_linear_ops = [
+            # Root decompositions given as RootLinearOperators.
+            self.create_linear_op()._linear_op,
+            # Cholesky (triangular) root decompositions of dense matrices.
+            MulLinearOperator(*(DenseLinearOperator(mat @ mat.mT + torch.eye(6)) for mat in mats)),
+        ]
+        positive = torch.rand(*batch_shape, 1, 1) + 0.5
+        mixed = positive.clone()
+        mixed.view(-1)[0] *= -1
+        for mul_linear_op, constant in itertools.product(mul_linear_ops, (positive, mixed)):
+            res = mul_linear_op * constant
+            # Positive constants keep the root decompositions.
+            self.assertEqual(isinstance(res, MulLinearOperator), constant is positive)
+            self.assertAllClose(res.to_dense(), mul_linear_op.to_dense() * constant)
+
+
 class TestMulLinearOperator(LinearOperatorTestCase, unittest.TestCase):
     seed = 10
 
@@ -32,7 +53,7 @@ def evaluate_linear_op(self, linear_op):
         return res
 
 
-class TestMulLinearOperatorBatch(LinearOperatorTestCase, unittest.TestCase):
+class TestMulLinearOperatorBatch(BatchConstantMulMixin, LinearOperatorTestCase, unittest.TestCase):
     seed = 2
 
     def create_linear_op(self):
@@ -51,7 +72,7 @@ def evaluate_linear_op(self, linear_op):
         return res
 
 
-class TestMulLinearOperatorMultiBatch(LinearOperatorTestCase, unittest.TestCase):
+class TestMulLinearOperatorMultiBatch(BatchConstantMulMixin, LinearOperatorTestCase, unittest.TestCase):
     seed = 1
     skip_slq_tests = True
 

test/operators/test_triangular_linear_operator.py

@@ -69,6 +69,18 @@ def test_solve_triangular(self):
         with self.assertRaisesRegex(RuntimeError, "Incompatible argument"):
             torch.linalg.solve_triangular(linear_op, rhs, upper=not is_upper)
 
+    def test_constant_mul_batch(self):
+        upper = self.create_linear_op().upper
+        for batch_shape in (torch.Size([5]), torch.Size([2, 3])):
+            tensor = torch.randn(*batch_shape, 5, 5)
+            tensor = tensor.triu() if upper else tensor.tril()
+            linear_op = TriangularLinearOperator(DenseLinearOperator(tensor), upper=upper)
+            constant = torch.rand(*batch_shape, 1, 1) + 0.5
+            res = linear_op * constant
+            self.assertIsInstance(res, TriangularLinearOperator)
+            self.assertEqual(res.upper, upper)
+            self.assertAllClose(res.to_dense(), tensor * constant)
+
     # Tests that we bypass because TriangularLinearOperators are not symmetric or PSD
 
     def test_add_low_rank(self):