From 8c3b994a3d1b010a84c65fd15fbcab047f9fb6c9 Mon Sep 17 00:00:00 2001 From: Anton Soloviev Date: Thu, 15 Oct 2026 12:53:24 -0700 Subject: [PATCH] Handle positive batch scales for low-rank-plus-diagonal operators --- ...ow_rank_root_added_diag_linear_operator.py | 2 +- test/operators/test_lowrank_batch_scaling.py | 35 +++++++++++++++++++ 2 files changed, 36 insertions(+), 1 deletion(-) create mode 100644 test/operators/test_lowrank_batch_scaling.py diff --git a/linear_operator/operators/low_rank_root_added_diag_linear_operator.py b/linear_operator/operators/low_rank_root_added_diag_linear_operator.py index 2453213c..d7e17bb8 100644 --- a/linear_operator/operators/low_rank_root_added_diag_linear_operator.py +++ b/linear_operator/operators/low_rank_root_added_diag_linear_operator.py @@ -50,7 +50,7 @@ def _mul_constant( self: LinearOperator, other: float | torch.Tensor # shape: (*batch, M, N) ) -> LinearOperator: # shape: (*batch, M, N) # We have to over-ride this here for the case where the constant is negative - if other > 0: + if (other > 0).all(): res = super()._mul_constant(other) else: res = AddedDiagLinearOperator(self._linear_op._mul_constant(other), self._diag_tensor._mul_constant(other)) diff --git a/test/operators/test_lowrank_batch_scaling.py b/test/operators/test_lowrank_batch_scaling.py new file mode 100644 index 00000000..c5aa5f12 --- /dev/null +++ b/test/operators/test_lowrank_batch_scaling.py @@ -0,0 +1,35 @@ +#!/usr/bin/env python3 + +import unittest + +import torch + +from linear_operator.operators import DiagLinearOperator, LowRankRootAddedDiagLinearOperator, LowRankRootLinearOperator + + +class TestLowRankBatchScaling(unittest.TestCase): + def test_batch_scale_values_and_gradients(self) -> None: + torch.manual_seed(6) + for batch in ((3,), (2, 3)): + with self.subTest(batch=batch): + root = torch.randn(*batch, 8, 3, dtype=torch.double, requires_grad=True) + diagonal = torch.ones(*batch, 8, dtype=torch.double, requires_grad=True) + scales = torch.linspace(0.5, 2.0, 3, dtype=torch.double, requires_grad=True) + copies = [value.detach().clone().requires_grad_() for value in (root, diagonal, scales)] + operator = LowRankRootLinearOperator(root) + DiagLinearOperator(diagonal) + scaled = operator * scales[:, None, None] + self.assertIsInstance(scaled, LowRankRootAddedDiagLinearOperator) + actual = scaled.to_dense() + dense = copies[0] @ copies[0].mT + torch.diag_embed(copies[1]) + expected = dense * copies[2][:, None, None] + torch.testing.assert_close(actual, expected) + actual_grads = torch.autograd.grad(actual.square().sum(), (root, diagonal, scales)) + expected_grads = torch.autograd.grad(expected.square().sum(), copies) + for actual_grad, expected_grad in zip(actual_grads, expected_grads): + torch.testing.assert_close(actual_grad, expected_grad) + + def test_batch_with_zero_factor(self) -> None: + root = torch.ones(3, 8, 2, dtype=torch.double) + operator = LowRankRootLinearOperator(root).add_diagonal(torch.ones(3, 8, dtype=torch.double)) + scales = torch.tensor([0.0, 0.5, 2.0], dtype=torch.double)[:, None, None] + torch.testing.assert_close((operator * scales).to_dense(), operator.to_dense() * scales)