Skip to content

Commit

Permalink
Update _linear_operator.py (#9)
Browse files Browse the repository at this point in the history
* Update _linear_operator.py

Add additional LinearOperator.__add__ case for Pythonic addition operations with the builtin sum() function. After other common use-cases, check if other is a numeric Python type and equal to zero, and if so, return self.

* Update linear_operator/operators/_linear_operator.py

Co-authored-by: Max Balandat <Balandat@users.noreply.github.com>

* Update _linear_operator.py

Co-authored-by: Geoff Pleiss <gpleiss@gmail.com>
Co-authored-by: Max Balandat <Balandat@users.noreply.github.com>
  • Loading branch information
3 people authored Sep 9, 2022
1 parent 8bc1d04 commit a58efe7
Showing 1 changed file with 4 additions and 3 deletions.
7 changes: 4 additions & 3 deletions linear_operator/operators/_linear_operator.py
Original file line number Diff line number Diff line change
Expand Up @@ -4,6 +4,7 @@

import functools
import math
import numbers
import warnings
from abc import ABC, abstractmethod
from copy import deepcopy
Expand Down Expand Up @@ -2546,9 +2547,7 @@ def zero_mean_mvn_samples(self, num_samples: int) -> torch.Tensor:

return samples

def __add__(self, other: Union[torch.Tensor, LinearOperator, float]) -> LinearOperator:
from torch import Tensor

def __add__(self, other: Union[Tensor, LinearOperator, float]) -> LinearOperator:
from .added_diag_linear_operator import AddedDiagLinearOperator
from .dense_linear_operator import to_linear_operator
from .diag_linear_operator import DiagLinearOperator
Expand All @@ -2568,6 +2567,8 @@ def __add__(self, other: Union[torch.Tensor, LinearOperator, float]) -> LinearOp
new_self = self if self.shape[:-2] == shape[:-2] else self._expand_batch(shape[:-2])
new_other = other if other.shape[:-2] == shape[:-2] else other._expand_batch(shape[:-2])
return SumLinearOperator(new_self, new_other)
elif isinstance(other, numbers.Number) and other == 0:
return self
else:
return SumLinearOperator(self, other)

Expand Down

0 comments on commit a58efe7

Please sign in to comment.