Allow negative indices in pick_nth_selected
#110
Conversation
|
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 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. |
penzai/core/selectors.py
Outdated
| pz.select((0, 3, -2)) \ | ||
| .at_instances_of(int) \ | ||
| .where(lambda i: i < 0) \ | ||
| .apply(lambda i: i + shift) |
There was a problem hiding this comment.
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
There was a problem hiding this comment.
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.)
There was a problem hiding this comment.
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
danieldjohnson
left a comment
danieldjohnson
left a comment
There was a problem hiding this comment.
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 for the full set of checks that gets run in CI.)
tests/core/selectors_test.py
Outdated
| 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 | ||
| ) |
There was a problem hiding this comment.
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.
penzai/core/selectors.py
Outdated
| pz.select((0, 3, -2)) \ | ||
| .at_instances_of(int) \ | ||
| .where(lambda i: i < 0) \ | ||
| .apply(lambda i: i + shift) |
There was a problem hiding this comment.
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.)
|
Thanks for the review! I'll address the comments some time this week. |
|
@danieldjohnson I've addressed your comments. All tests pass locally. Could you have another look? |
2 participants
Resolves #109