Merge pull request #70 from Turakar/smallopt_repeat

Add short path for LinearOperator.repeat()

Author: Balandat

linear_operator/operators/_linear_operator.py

@@ -2023,6 +2023,10 @@ def repeat(self, *sizes: Union[int, Tuple[int, ...]]) -> LinearOperator:
         """
         from .batch_repeat_linear_operator import BatchRepeatLinearOperator
 
+        # Short path if no repetition is necessary
+        if all(x == 1 for x in sizes) and len(sizes) == self.dim():
+            return self
+
         if len(sizes) < 3 or tuple(sizes[-2:]) != (1, 1):
             raise RuntimeError(
                 "Invalid repeat arguments {}. Currently, repeat only works to create repeated "

test/utils/test_repeat.py

@@ -0,0 +1,23 @@
+import unittest
+
+import torch
+
+from linear_operator.operators.dense_linear_operator import DenseLinearOperator
+from linear_operator.test.utils import approx_equal
+
+
+class TestRepeat(unittest.TestCase):
+    def make_example(self):
+        return DenseLinearOperator(torch.randn(3, 3))
+
+    def test_repeat(self):
+        example = self.make_example()
+        repeated = example.repeat(2, 1, 1)
+        repeated_dense = example.to_dense().repeat(2, 1, 1)
+        self.assertTrue(approx_equal(repeated.to_dense(), repeated_dense))
+
+    def test_repeat_noop(self):
+        example = self.make_example()
+        repeated = example.repeat(1, 1)
+        self.assertTrue(approx_equal(repeated.to_dense(), example.to_dense()))
+        self.assertIsInstance(repeated, DenseLinearOperator)  # ensure that fast path is taken