Merge branch 'main' into main

linear_operator/operators/_linear_operator.py

@@ -1752,7 +1752,7 @@ def mul(self, other: Union[float, torch.Tensor, "LinearOperator"]) -> LinearOper
             other = torch.tensor(other, dtype=self.dtype, device=self.device)
 
         try:
-            torch.broadcast_shapes(self.shape, other.shape)
+            broadcast_shape = torch.broadcast_shapes(self.shape, other.shape)
         except RuntimeError:
             raise RuntimeError(
                 "Cannot multiply LinearOperator of size {} by an object of size {}".format(self.shape, other.shape)
@@ -1761,7 +1761,7 @@ def mul(self, other: Union[float, torch.Tensor, "LinearOperator"]) -> LinearOper
         if torch.is_tensor(other):
             if other.numel() == 1:
                 return self._mul_constant(other.squeeze())
-            elif other.shape[-2:] == torch.Size((1, 1)):
+            elif other.shape[-2:] == torch.Size((1, 1)) and self.batch_shape == broadcast_shape[:-2]:
                 return self._mul_constant(other.view(*other.shape[:-2]))
 
         return self._mul_matrix(to_linear_operator(other))

linear_operator/operators/batch_repeat_linear_operator.py

@@ -21,7 +21,7 @@ def __init__(self, base_linear_op, batch_repeat=torch.Size((1,))):
                 )
             if isinstance(base_linear_op, BatchRepeatLinearOperator):
                 raise RuntimeError(
-                    "BatchRepeatLinearOperator recieved the following args:\n"
+                    "BatchRepeatLinearOperator received the following args:\n"
                     "base_linear_op: {} (size: {}), batch_repeat: {}.".format(
                         base_linear_op, base_linear_op.shape, batch_repeat
                     )

test/operators/test_constant_mul_linear_operator.py

@@ -4,7 +4,7 @@
 
 import torch
 
-from linear_operator.operators import ToeplitzLinearOperator
+from linear_operator.operators import DenseLinearOperator, ToeplitzLinearOperator
 from linear_operator.test.linear_operator_test_case import LinearOperatorTestCase
 from linear_operator.utils.toeplitz import sym_toeplitz
 
@@ -42,7 +42,7 @@ def evaluate_linear_op(self, linear_op):
 
 class TestConstantMulLinearOperatorMultiBatch(LinearOperatorTestCase, unittest.TestCase):
     seed = 0
-    # Because these LTs are large, we'll skil the big tests
+    # Because these LTs are large, we'll skip the big tests
     should_test_sample = False
     skip_slq_tests = True
 
@@ -60,7 +60,7 @@ def evaluate_linear_op(self, linear_op):
 
 class TestConstantMulLinearOperatorMultiBatchBroadcastConstant(LinearOperatorTestCase, unittest.TestCase):
     seed = 0
-    # Because these LTs are large, we'll skil the big tests
+    # Because these LTs are large, we'll skip the big tests
     should_test_sample = False
     skip_slq_tests = True
 
@@ -76,5 +76,22 @@ def evaluate_linear_op(self, linear_op):
         return toeplitz.to_dense() * constant
 
 
+class TestConstantMulLinearOperatorBatchBroadcastOperator(LinearOperatorTestCase, unittest.TestCase):
+    """Test which broadcasts the operator to match the constant tensor's batch size, see Github issue #33"""
+
+    seed = 0
+    should_test_sample = False
+    skip_slq_tests = True
+
+    def create_linear_op(self):
+        mat = torch.randn(5, 6)
+        mat = mat.matmul(mat.mT).reshape(1, 5, 5)
+        constant = torch.randn(2, 1, 1).abs()
+        return DenseLinearOperator(mat) * constant
+
+    def evaluate_linear_op(self, linear_op):
+        return linear_op.to_dense()
+
+
 if __name__ == "__main__":
     unittest.main()