From a2aa18cf89193c8f9de39039a8932a15437f3b99 Mon Sep 17 00:00:00 2001 From: Mark Colley <39366189+M-Colley@users.noreply.github.com> Date: Sat, 17 Oct 2026 04:29:35 +0100 Subject: [PATCH] Support batches of constants in MulLinearOperator._mul_constant `LinearOperator.mul` passes a `*batch`-shaped tensor of constants to `_mul_constant` when multiplying by a `*batch x 1 x 1` tensor. - `MulLinearOperator._mul_constant` used `if other > 0`, which raises "Boolean value of Tensor with more than one value is ambiguous" for more than one constant. It now uses `(other > 0).all()`, as `RootLinearOperator._mul_constant` does. - `TriangularLinearOperator._mul_constant` multiplied its `*batch x N x N` tensor by `other.unsqueeze(-1)`, which scales the rows instead of the batch members (a wrong result if the batch size equals N, an error otherwise). It is reached from the fixed branch above through the Cholesky roots of the `RootLinearOperator`s. It now uses `other.unsqueeze(-1).unsqueeze(-1)`. Co-Authored-By: Claude Opus 5.5 --- .../operators/mul_linear_operator.py | 2 +- .../operators/triangular_linear_operator.py | 2 +- test/operators/test_mul_linear_operator.py | 27 ++++++++++++++++--- .../test_triangular_linear_operator.py | 12 +++++++++ 4 files changed, 38 insertions(+), 5 deletions(-) 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):