Repository: google-deepmind/penzai
Pull request #110: Allow negative indices in `pick_nth_selected`
Author: login: JEM-Mosig
Created: 2025-04-14T08:42:48Z | Merged: 2025-04-26T07:06:25Z

Resolves https://github.com/google-deepmind/penzai/issues/109


Diff:
diff --git a/penzai/core/selectors.py b/penzai/core/selectors.py
index d89f8e0..3832fd9 100644
--- a/penzai/core/selectors.py
+++ b/penzai/core/selectors.py
@@ -39,6 +39,42 @@
 T = typing.TypeVar("T")
 
 
+def _shift_negative_indices(
+    indices: Iterable[int], shift: int
+) -> tuple[int, ...]:
+  """Adds `shift` to negative indices and leaves non-negative indices unchanged
+
+  Can be used to handle negative indices. For example, if we expect indices in
+  `r = range(6)` and we get `[0, 3, -2]` as input, we can use ::
+
+    shift_negative_indices([0, 3, -2], len(r))
+
+  to get `(0, 3, 4)`. The same can be achieved in more generality with ::
+
+    pz.select((0, 3, -2))
+      .at_instances_of(int)
+      .where(lambda i: i < 0)
+      .apply(lambda i: i + shift)
+
+  which is why this method is private to this module
+
+  Args:
+    indices: The integers to shift
+    shift: The offset to add to negative indices. Usually, this is the largest
+      index + 1, i.e. the length of the range of indices
+
+  Returns:
+    The indices as a tuple, with negative indices increased by `shift`
+  """
+  maybe_shifted_indices = []
+  for index in indices:
+    if index < 0:
+      maybe_shifted_indices.append(index + shift)
+    else:
+      maybe_shifted_indices.append(index)
+  return tuple(maybe_shifted_indices)
+
+
 @struct.pytree_dataclass
 class SelectionHole(struct.Struct):
   """A hole in a structure, taking the place of a selected subtree.
@@ -1356,6 +1392,8 @@ def pick_nth_selected(self, n: int | Sequence[int]) -> Selection:
     else:
       indices = n
 
+    indices = _shift_negative_indices(indices, len(self.selected_by_path))
+
     with _wrap_selection_errors(self):
       keep = _InProgressSelectionBoundary
       new_selected_by_path = collections.OrderedDict({
diff --git a/tests/core/selectors_test.py b/tests/core/selectors_test.py
index 66eb1a2..f296ba0 100644
--- a/tests/core/selectors_test.py
+++ b/tests/core/selectors_test.py
@@ -56,6 +56,18 @@ class SELECTED_PART:  # pylint: disable=invalid-name
 
 class SelectorsTest(absltest.TestCase):
 
+  def test_shift_negative_indices(self):
+    for input_indices, shift, expected_output in [
+        ((), 1, ()),
+        ([0, 3, -2], len(range(6)), (0, 3, 4)),
+    ]:
+      self.assertEqual(
+          penzai.core.selectors._shift_negative_indices(
+              input_indices, shift=shift
+          ),
+          expected_output,
+      )
+
   def test_select(self):
     self.assertEqual(
         pz.select(make_example_object()),
@@ -565,6 +577,26 @@ def test_pick_nth_selected(self):
         ),
         [0, 1, 2, 3, SELECTED_PART(value=4), 5, 6, 7, 8, 9],
     )
+    # Test negative indices for `pick_nth_selected`
+    self.assertEqual(
+        (
+            pz.select(list(range(10)))
+            .at_instances_of(int)
+            .pick_nth_selected(-2)
+            .apply(SELECTED_PART)
+        ),
+        [0, 1, 2, 3, 4, 5, 6, 7, SELECTED_PART(value=8), 9],
+    )
+    # Don't select anything if index is out of range
+    self.assertEqual(
+        (
+            pz.select([0, 1, 2])
+            .at_instances_of(int)
+            .pick_nth_selected(5)
+            .apply(SELECTED_PART)
+        ),
+        [0, 1, 2],
+    )
 
   def test_invert__example_1(self):
     predicate = lambda node: isinstance(node, CustomLeaf) and node.tag <= 10


Reviews:
user: login: JEM-Mosig
state: COMMENTED
author_association: CONTRIBUTOR
submitted_at: 2025-04-15T02:56:36Z
user: login: danieldjohnson
body: Thanks for the PR!

It looks like one of our CI steps is currently failing (`uv run pyink penzai tests --check`). Mind fixing this? If you run `uv run pyink penzai tests` it should reformat everything automatically.

(See [here](https://github.com/google-deepmind/penzai/blob/main/.github/workflows/unittests.yml) for the full set of checks that gets run in CI.)
state: COMMENTED
author_association: COLLABORATOR
submitted_at: 2025-04-20T21:56:24Z
user: login: JEM-Mosig
state: COMMENTED
author_association: CONTRIBUTOR
submitted_at: 2025-04-24T07:30:24Z
user: login: danieldjohnson
body: Thanks much!
state: APPROVED
author_association: COLLABORATOR
submitted_at: 2025-04-26T06:44:37Z

Review comments:
diff_hunk: @@ -39,6 +39,44 @@
 T = typing.TypeVar("T")
 
 
+def _shift_negative_indices(indices: Iterable[int], shift: int) -> tuple[int, ...]:
+  """Adds `shift` to negative indices and leaves non-negative indices unchanged
+
+  Can be used to handle negative indices. For example, if we expect indices in
+  `r = range(6)` and we get `[0, 3, -2]` as input, we can use
+
+  ```py
+  shift_negative_indices([0, 3, -2], len(r))
+  ```
+
+  to get `(0, 3, 4)`. The same can be achieved in more generality with
+
+  ```py
+  pz.select((0, 3, -2)) \
+    .at_instances_of(int) \
+    .where(lambda i: i < 0) \
+    .apply(lambda i: i + shift)
path: penzai/core/selectors.py
user: login: JEM-Mosig
body: Not sure if I should just use this instead of defining `_shift_negative_indices`. I'm weary of doing so because we're in the module that defines selectors here. Also, we don't need that kind of generality at this point and `_shift_negative_indices` avoids some overhead
created_at: 2025-04-15T02:56:36Z
author_association: CONTRIBUTOR
diff_hunk: @@ -20,12 +20,31 @@
 
 import collections
 import dataclasses
-from typing import Any
+from typing import Any, Iterable
 
 from absl.testing import absltest
 import jax
 from penzai import pz
 import penzai.core.selectors
+import pytest
+
+
+@pytest.mark.parametrize(
+  "input_indices, shift, expected_output",
+  [
+    ((), 1, ()),
+    ([0, 3, -2], len(range(6)), (0, 3, 4)),
+  ]
+)
+def test_shift_negative_indices(
+  input_indices: Iterable[int],
+  shift: int,
+  expected_output: tuple[int, ...],
+):
+  assert (
+    penzai.core.selectors._shift_negative_indices(input_indices, shift=shift)
+    == expected_output
+  )
path: tests/core/selectors_test.py
user: login: danieldjohnson
body: Would you mind adding this as an absltest-style test (as a method inside SelectorsTest) instead of using pytest here? The rest of the tests do not use pytest yet and I'd prefer to be consistent here for now.
created_at: 2025-04-20T21:48:58Z
author_association: COLLABORATOR
diff_hunk: @@ -39,6 +39,44 @@
 T = typing.TypeVar("T")
 
 
+def _shift_negative_indices(indices: Iterable[int], shift: int) -> tuple[int, ...]:
+  """Adds `shift` to negative indices and leaves non-negative indices unchanged
+
+  Can be used to handle negative indices. For example, if we expect indices in
+  `r = range(6)` and we get `[0, 3, -2]` as input, we can use
+
+  ```py
+  shift_negative_indices([0, 3, -2], len(r))
+  ```
+
+  to get `(0, 3, 4)`. The same can be achieved in more generality with
+
+  ```py
+  pz.select((0, 3, -2)) \
+    .at_instances_of(int) \
+    .where(lambda i: i < 0) \
+    .apply(lambda i: i + shift)
path: penzai/core/selectors.py
user: login: danieldjohnson
body: Defining a new function here like you've done makes sense to me, the logic is probably too simple to benefit much from using selectors. (You could also just inline the logic into the call site since it's only used once.)
created_at: 2025-04-20T21:56:10Z
author_association: COLLABORATOR
diff_hunk: @@ -39,6 +39,44 @@
 T = typing.TypeVar("T")
 
 
+def _shift_negative_indices(indices: Iterable[int], shift: int) -> tuple[int, ...]:
+  """Adds `shift` to negative indices and leaves non-negative indices unchanged
+
+  Can be used to handle negative indices. For example, if we expect indices in
+  `r = range(6)` and we get `[0, 3, -2]` as input, we can use
+
+  ```py
+  shift_negative_indices([0, 3, -2], len(r))
+  ```
+
+  to get `(0, 3, 4)`. The same can be achieved in more generality with
+
+  ```py
+  pz.select((0, 3, -2)) \
+    .at_instances_of(int) \
+    .where(lambda i: i < 0) \
+    .apply(lambda i: i + shift)
path: penzai/core/selectors.py
user: login: JEM-Mosig
body: > You could also just inline the logic into the call site since it's only used once

I'd prefer to keep it separate so code blocks stay short and hence easy to interpret
created_at: 2025-04-24T07:30:24Z
author_association: CONTRIBUTOR

PR comments:
user: login: google-cla[bot]
created_at: 2025-04-14T08:42:53Z
body: Thanks for your pull request! It looks like this may be your first contribution to a Google open source project. Before we can look at your pull request, you'll need to sign a Contributor License Agreement (CLA).

View this [failed invocation](https://github.com/google-deepmind/penzai/pull/110/checks?check_run_id=40491876012) of the CLA check for more information.

For the most up to date status, view the checks section at the bottom of the pull request.
author_association: NONE
user: login: JEM-Mosig
created_at: 2025-04-22T08:12:56Z
body: Thanks for the review! I'll address the comments some time this week.
author_association: CONTRIBUTOR
user: login: JEM-Mosig
created_at: 2025-04-24T07:48:08Z
body: @danieldjohnson I've addressed your comments. All tests pass locally. Could you have another look?
author_association: CONTRIBUTOR