Refactor advanced subtensor #1756

jaanerik · 2025-11-28T14:30:51Z

Description

Allows vectorizing AdvancedSetSubtensor.

Gemini picks up where Copilot left off.

Related Issue

Closes Reconsider use of SliceType and NoneType Variables as inputs to AdvancedIndexing #541
Directly follows from Refactor AdvancedSubtensor to only take numerical inputs, unify newaxis handling, and update all dispatch functions #1622 and related to pymc-extras#577

Checklist

Checked that the pre-commit linting/style checks pass (ruff removed whitespace etc from copilot commits also I think)
Changed tests that prove the fix is effective or that the new feature works
Added necessary documentation (docstrings and/or example notebooks)
If you are a pro: each commit corresponds to a relevant logical change

Type of change

ricardoV94

This is looking pretty good.

Mostly I want to believe there is still room to simplify things / reuse code.
This is also a good opportunity to simplify the idx_list. There's no reason to use ScalarTypes in the dummy slices, and it's complicating our equality and hashing.

What about using simple integers to indicate what is the role of each index variable?

old_idx_list = (ps.int64, slice(ps.int64, None, None), ps.int64, slice(ps.int64, None, ps.int64))
new_idx_list = (0, slice(1, None, None), 2, slice(3, None, 4))

Having the indices could probably come in handy anyway. With this we shouldn't need a custom hash / eq, we can just use the default one from __props__.

pytensor/tensor/rewriting/subtensor.py

ricardoV94 · 2025-12-09T14:10:54Z

pytensor/tensor/rewriting/subtensor.py

    else:
-        x, y, *idxs = node.inputs
+        x, y = node.inputs[0], node.inputs[1]
+        tensor_inputs = node.inputs[2:]


I don't like the name tensor_inputs, x, y are also tensor and inputs. Use index_variables?

This applies elsewhere

Using index_variables now.

Still using tensor_inputs in the last code you pushed

Sorry about that, renamed now.

pytensor/tensor/rewriting/subtensor.py

ricardoV94 · 2025-12-09T14:14:30Z

pytensor/tensor/rewriting/subtensor.py

        # must already be raveled in the original graph, so we don't need to do anything to it
-        new_out = node.op(raveled_x, y, *new_idxs)
-        # But we must reshape the output to math the original shape
+        new_out = AdvancedIncSubtensor(


You should use type(op) so that subclasses are respected. It may also make sense to add a method to these indexing Ops like op.with_new_indices() that clones itself with a new idx_list. Maybe that will be the one that handles creating the new idx_list, instead of having to be here in the rewrite.

Currently functions ravel_multidimensional_bool_idx and ravel_multidimensional_bool_idx don't assume the subclass in the same way, but it'd be nice if you could check. Also, if I am giving up on some rewrites too quickly here, please let me know.

ricardoV94 · 2025-12-09T14:16:59Z

pytensor/tensor/subtensor.py

+    def __init__(self, idx_list):
+        """
+        Initialize AdvancedSubtensor with index list.
+
+        Parameters
+        ----------
+        idx_list : tuple
+            List of indices where slices are stored as-is,
+            and numerical indices are replaced by their types.
+        """
+        self.idx_list = tuple(
+            index_vars_to_types(idx, allow_advanced=True) for idx in idx_list
+        )
+        # Store expected number of tensor inputs for validation
+        self.expected_inputs_len = len(
+            get_slice_elements(self.idx_list, lambda entry: isinstance(entry, Type))
+        )
+
+    def __hash__(self):
+        msg = []
+        for entry in self.idx_list:
+            if isinstance(entry, slice):
+                msg += [(entry.start, entry.stop, entry.step)]
+            else:
+                msg += [entry]
+
+        idx_list = tuple(msg)
+        return hash((type(self), idx_list))


This already exists in Subtensor? If so create a BaseSubtensor class that handles idx_list and hash/equality based on it.

Make all Subtensor operations inherit from it

Note this advice may make no sense if we simplify the idx_list to not need custom hash / eq

hash is a bit different based on whether it is *IncSubtensor or not in the current implementation. Wrote about the Python 3.11 slice not being hashable below.

ricardoV94 · 2025-12-09T14:18:33Z

pytensor/tensor/subtensor.py

-            )
-        else:
-            return vectorize_node_fallback(op, node, batch_x, *batch_idxs)
+        # With the new interface, all inputs are tensors, so Blockwise can handle them


Comment should not mention a specific time period. Previous status is not relevant here

Also we still want to avoid Blockwise eagerly if we can

All time periods should be removed from comments that were added in this PR.

pytensor/tensor/variable.py

ricardoV94 · 2025-12-09T14:35:27Z

pytensor/tensor/variable.py

+                pattern.append("x")
+                new_args.append(slice(None))
+            else:
+                # Check for boolean index which consumes multiple dimensions


This is probably right, but why does it matter that boolean indexing consumes multiple dimensions? Aren't we doing expand_dims where there was None -> replace new_axis by None slice -> index again?

There might be a more dynamic approach, but I have had trouble with multidimensional bool arrays and ellipsis, e.g. x[multi_dim_bool_tensor,None,...] will not know where to add the new axis.

tests/tensor/test_subtensor.py

ricardoV94 · 2025-12-09T14:40:35Z

Just wanted to repeat, this is looking great. Thanks so far @jaanerik

I'm being picky because indexing is a pretty fundamental operation, so want to make sure we get it right this time.

jaanerik · 2025-12-14T17:17:35Z

This is looking pretty good.

Mostly I want to believe there is still room to simplify things / reuse code. This is also a good opportunity to simplify the idx_list. There's no reason to use ScalarTypes in the dummy slices, and it's complicating our equality and hashing.

What about using simple integers to indicate what is the role of each index variable?
old_idx_list = (ps.int64, slice(ps.int64, None, Non), ps.int64, slice(3, ps.int64, 4))
new_idx_list = (0, slice(1, None, None), 2, slice(3, None, 4))
Having the indices could probably come in handy anyway. With this we shouldn't need a custom hash / eq, we can just use the default one from __props__.

@ricardoV94 I am struggling a bit with understanding your example, because old_idx_list also has ints. Could you clarify how you see constant index and variable index both working here.

If you have the last slice of old_idx_list as slice(3, ps.int64, 4) and new one has slice(3, 4, 5) (this does not correspond exactly to your example, because I didn't understand it) then would 3 from new_idx_list indicate a Constant(3), 4 would indicate a Scalar() and 5 would indicate Constant(4)? Or would you discriminate between a static index and a variable index in a different way?

Absolutely no problem with being picky. I am very grateful for the feedback :)

ricardoV94 · 2025-12-14T17:47:07Z

I updated it, it was some copy-paste typos. Old_idx doesn't have ints, only ps.int64 to mark the place in the pattern that have input variables.

pytensor/link/numba/dispatch/subtensor.py

pytensor/tensor/conv/abstract_conv.py

pytensor/tensor/rewriting/subtensor.py

jaanerik · 2026-01-06T22:37:01Z

Taking a while with this, sorry. Finally simplified the idx_list.

I'll note that Python slice became hashable in Python 3.12 GH-101264. If pyproject.toml supports Python 3.11 still, then we have to introduce the explicit __hash__ in each subtensor class in pytensor/tensor/subtensor.py that uses _hashable_idx_list() prop, I think.

Co-authored-by: ricardoV94 <[email protected]>

… handling Co-authored-by: ricardoV94 <[email protected]>

…ation Co-authored-by: ricardoV94 <[email protected]>

…ensor approach Co-authored-by: ricardoV94 <[email protected]>

… interface, store expected_inputs_len Co-authored-by: ricardoV94 <[email protected]>

jaanerik force-pushed the refactor-advanced-subtensor branch 2 times, most recently from 3ce54c7 to 92f61ed Compare December 1, 2025 08:23

jaanerik marked this pull request as ready for review December 2, 2025 09:49

jaanerik force-pushed the refactor-advanced-subtensor branch from 92f61ed to a6cb68d Compare December 2, 2025 11:23

jaanerik mentioned this pull request Dec 2, 2025

Implement AdvancedIncSubtensorExplicit #1751

Closed

11 tasks

jaanerik force-pushed the refactor-advanced-subtensor branch 3 times, most recently from 546100c to 4b02064 Compare December 9, 2025 12:21

ricardoV94 requested changes Dec 9, 2025

View reviewed changes

ricardoV94 reviewed Dec 9, 2025

View reviewed changes

pytensor/tensor/variable.py Outdated Show resolved Hide resolved

ricardoV94 reviewed Dec 9, 2025

View reviewed changes

pytensor/tensor/variable.py Outdated Show resolved Hide resolved

ricardoV94 reviewed Dec 9, 2025

View reviewed changes

tests/tensor/test_subtensor.py Outdated Show resolved Hide resolved

ricardoV94 reviewed Dec 9, 2025

View reviewed changes

tests/tensor/test_subtensor.py Show resolved Hide resolved

ricardoV94 mentioned this pull request Dec 9, 2025

Refactor AdvancedSubtensor to only take numerical inputs, unify newaxis handling, and update all dispatch functions #1622

Closed

ricardoV94 added major Op implementation indexing labels Dec 9, 2025

jaanerik force-pushed the refactor-advanced-subtensor branch 4 times, most recently from 4787bc9 to 5ba887b Compare December 30, 2025 00:41

ricardoV94 reviewed Dec 30, 2025

View reviewed changes

pytensor/link/numba/dispatch/subtensor.py Outdated Show resolved Hide resolved

ricardoV94 reviewed Dec 30, 2025

View reviewed changes

pytensor/tensor/conv/abstract_conv.py Outdated Show resolved Hide resolved

ricardoV94 reviewed Dec 30, 2025

View reviewed changes

pytensor/tensor/rewriting/subtensor.py Outdated Show resolved Hide resolved

jaanerik force-pushed the refactor-advanced-subtensor branch 2 times, most recently from d821794 to 02a8d24 Compare January 6, 2026 13:37

jaanerik force-pushed the refactor-advanced-subtensor branch 3 times, most recently from 71aeef3 to b5d5ffa Compare January 6, 2026 22:35

jaanerik force-pushed the refactor-advanced-subtensor branch 2 times, most recently from e2af59e to eafb683 Compare January 7, 2026 10:04

Copilot AI and others added 8 commits January 7, 2026 17:18

Initial plan

71d5e35

Implement core refactoring of AdvancedSubtensor and AdvancedIncSubtensor

5850fcc

Co-authored-by: ricardoV94 <[email protected]>

Complete refactoring with improved factory functions and proper slice…

c3a81e2

… handling Co-authored-by: ricardoV94 <[email protected]>

Final fix: use as_index_variable consistently with original implement…

006b738

…ation Co-authored-by: ricardoV94 <[email protected]>

Refactor newaxis handling: move to __getitem__ level, unify with Subt…

fe1188b

…ensor approach Co-authored-by: ricardoV94 <[email protected]>

Update dispatch functions and rewrite rules for new AdvancedSubtensor…

97ed8ab

… interface, store expected_inputs_len Co-authored-by: ricardoV94 <[email protected]>

Finish Copilot code

5d47e62

Replace np.newaxis with None, remove NoneConst from indexing

9148eb7

jaanerik force-pushed the refactor-advanced-subtensor branch from eafb683 to 0be86da Compare January 8, 2026 14:25

jaanerik added 10 commits January 8, 2026 16:57

Fix rewriting, use existing functions, respect subclasses

aee4873

Fix tests

a88c68b

Implement BaseSubtensor

58a932c

Fix rebase

3f77d83

Remove one test that incorrectly extends materialized Ops

5f60074

Rename tensor_inputs to index_variables

f94da4a

Fix JAX dispatch

eba99ae

Refactor AdvancedSubtensor1 to use idx_list, add comments

525003c

Implement simpler idx_list

a2065c2

Revert unrelated code, remove deprecated code from a test

e89da66

jaanerik force-pushed the refactor-advanced-subtensor branch from 0be86da to e89da66 Compare January 8, 2026 14:59

Refactor advanced subtensor #1756

Are you sure you want to change the base?

Refactor advanced subtensor #1756

Uh oh!

Conversation

jaanerik commented Nov 28, 2025

Description

Related Issue

Checklist

Type of change

Uh oh!

ricardoV94 left a comment • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

ricardoV94 Dec 30, 2025 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

jaanerik Jan 7, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Uh oh!

ricardoV94 commented Dec 9, 2025 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

jaanerik commented Dec 14, 2025

Uh oh!

ricardoV94 commented Dec 14, 2025 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

jaanerik commented Jan 6, 2026

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

2 participants

ricardoV94 left a comment •

edited

Loading

ricardoV94 Dec 30, 2025 •

edited

Loading

jaanerik Jan 7, 2026 •

edited

Loading

ricardoV94 commented Dec 9, 2025 •

edited

Loading

ricardoV94 commented Dec 14, 2025 •

edited

Loading