diff --git a/linear_operator/operators/mul_linear_operator.py b/linear_operator/operators/mul_linear_operator.py index e7095575..a4f7857a 100644 --- a/linear_operator/operators/mul_linear_operator.py +++ b/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 diff --git a/linear_operator/operators/triangular_linear_operator.py b/linear_operator/operators/triangular_linear_operator.py index fb5828b3..ff2a8a82 100644 --- a/linear_operator/operators/triangular_linear_operator.py +++ b/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) diff --git a/test/operators/test_mul_linear_operator.py b/test/operators/test_mul_linear_operator.py index 9c8b2d2a..a1a49951 100644 --- a/test/operators/test_mul_linear_operator.py +++ b/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 diff --git a/test/operators/test_triangular_linear_operator.py b/test/operators/test_triangular_linear_operator.py index 3bac5649..d7517b91 100644 --- a/test/operators/test_triangular_linear_operator.py +++ b/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):