Skip to content

Commit

Permalink
Fix issue with negative indexing in CatLinearOperator (#80)
Browse files Browse the repository at this point in the history
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.
  • Loading branch information
Balandat committed Sep 20, 2023
1 parent 0b51316 commit 5496242
Showing 1 changed file with 4 additions and 2 deletions.
6 changes: 4 additions & 2 deletions linear_operator/operators/cat_linear_operator.py
Original file line number Diff line number Diff line change
Expand Up @@ -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()
Expand Down

0 comments on commit 5496242

Please sign in to comment.