From 54962429ab89e2a9e519de6da8853513236b283b Mon Sep 17 00:00:00 2001 From: Max Balandat Date: Tue, 19 Sep 2023 17:29:56 -0700 Subject: [PATCH] Fix issue with negative indexing in `CatLinearOperator` (#80) Addresses #79 Looks like there may be a number of places where negative indexing isn't properly supported. This should probably be audited more comprehensively. --- linear_operator/operators/cat_linear_operator.py | 6 ++++-- 1 file changed, 4 insertions(+), 2 deletions(-) diff --git a/linear_operator/operators/cat_linear_operator.py b/linear_operator/operators/cat_linear_operator.py index a11b951d..1adc300e 100644 --- a/linear_operator/operators/cat_linear_operator.py +++ b/linear_operator/operators/cat_linear_operator.py @@ -111,8 +111,10 @@ def _split_slice(self, slice_idx: slice) -> Tuple[Sequence[int], List[slice]]: if slice_idx.step is not None: # TODO: Add support for this eventually. raise RuntimeError("Slicing a CatLinearOperator with a step is not currently supported!") - start_idx = slice_idx.start if slice_idx.start is not None else 0 - stop_idx = slice_idx.stop if slice_idx.stop is not None else self.size(self.cat_dim) + + cat_size = self.size(self.cat_dim) + start_idx = slice_idx.start % cat_size if slice_idx.start is not None else 0 + stop_idx = slice_idx.stop % cat_size if slice_idx.stop is not None else cat_size first_tensor_idx = self.idx_to_tensor_idx[start_idx].item() last_tensor_idx = self.idx_to_tensor_idx[stop_idx - 1].item()