diff --git a/linear_operator/operators/diag_linear_operator.py b/linear_operator/operators/diag_linear_operator.py index 3fcc576f..2fec3e59 100644 --- a/linear_operator/operators/diag_linear_operator.py +++ b/linear_operator/operators/diag_linear_operator.py @@ -369,7 +369,7 @@ def _mul_matrix( "Dimension Mismatch: Must have same diag_shape, but got " f"{self.diag_shape} and {other.diag_shape}" ) - return self.__class__(self.diag_values * other.diag_values, diag_shape=self.diag_shape) + return ConstantDiagLinearOperator(self.diag_values * other.diag_values, diag_shape=self.diag_shape) return super()._mul_matrix(other) def _prod_batch(self, dim: int) -> LinearOperator: diff --git a/linear_operator/operators/identity_linear_operator.py b/linear_operator/operators/identity_linear_operator.py index 54bf413e..f8fd686a 100644 --- a/linear_operator/operators/identity_linear_operator.py +++ b/linear_operator/operators/identity_linear_operator.py @@ -103,12 +103,6 @@ def _mul_constant( ) -> LinearOperator: # shape: (*batch, M, N) return ConstantDiagLinearOperator(self.diag_values * other, diag_shape=self.diag_shape) - def _mul_matrix( - self: LinearOperator, # shape: (..., #M, #N) - other: torch.Tensor | LinearOperator, # shape: (..., #M, #N) - ) -> LinearOperator: # shape: (..., M, N) - return other - def _permute_batch(self, *dims: int) -> LinearOperator: batch_shape = self.diag_values.permute(*dims, -1).shape[:-1] return IdentityLinearOperator( diff --git a/test/operators/test_identity_linear_operator.py b/test/operators/test_identity_linear_operator.py index ed9792a3..e51e0cbd 100644 --- a/test/operators/test_identity_linear_operator.py +++ b/test/operators/test_identity_linear_operator.py @@ -1,10 +1,11 @@ #!/usr/bin/env python3 +import math import unittest import torch -from linear_operator.operators import IdentityLinearOperator, to_dense +from linear_operator.operators import ConstantDiagLinearOperator, DenseLinearOperator, IdentityLinearOperator, to_dense from linear_operator.test.linear_operator_test_case import LinearOperatorTestCase @@ -107,6 +108,31 @@ def test_logdet(self): linear_op = self.create_linear_op() self.assertAllClose(torch.logdet(linear_op), torch.zeros(linear_op.batch_shape)) + def test_mul_matrix(self): + linear_op = self.create_linear_op() + other = torch.randn(*linear_op.shape) + expected = self.evaluate_linear_op(linear_op) * other + for rhs in (other, DenseLinearOperator(other)): + self.assertAllClose((linear_op * rhs).to_dense(), expected) + self.assertAllClose(torch.mul(linear_op, rhs).to_dense(), expected) + + def test_mul_constant_diag(self): + linear_op = self.create_linear_op() + others = ( + self.create_linear_op(), + IdentityLinearOperator(5), + ConstantDiagLinearOperator(torch.tensor([2.5]), diag_shape=5), + ConstantDiagLinearOperator( + torch.linspace(0.5, 1.5, math.prod(linear_op.batch_shape)).view(*linear_op.batch_shape, 1), + diag_shape=5, + ), + ) + for other in others: + expected = self.evaluate_linear_op(linear_op) * other.to_dense() + self.assertAllClose((linear_op * other).to_dense(), expected) + self.assertAllClose((other * linear_op).to_dense(), expected) + self.assertAllClose(torch.mul(linear_op, other).to_dense(), expected) + def test_sqrt_inv_matmul(self): linear_op = self.create_linear_op().detach().requires_grad_(True) if len(linear_op.batch_shape):