Refactor advanced subtensor by jaanerik · Pull Request #1756 · pymc-devs/pytensor

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

ricardoV94

This is shaping really great. I left some comments that may simplify things further

pytensor/link/numba/dispatch/subtensor.py

ricardoV94 · 2026-01-30T11:35:56Z

pytensor/link/numba/dispatch/subtensor.py

    if isinstance(op, AdvancedSubtensor1 | AdvancedSubtensor):
-        x, *idxs = node.inputs
+        x = node.inputs[0]
+        if isinstance(op, AdvancedSubtensor):
+            index_variables = node.inputs[1:]
+        else:
+            index_variables = node.inputs[1:]
    else:
-        x, y, *idxs = node.inputs
+        x, y = node.inputs[:2]
+        index_variables = node.inputs[2:]


Bot keeps trying to split this fine tuple unpacking into multiple lines, please revert

pytensor/link/numba/dispatch/subtensor.py

pytensor/tensor/rewriting/subtensor.py

ricardoV94 · 2026-01-30T11:45:09Z

pytensor/tensor/rewriting/subtensor.py

+
+        if (
+            len(inc_inputs) == len(sub_inputs)
+            and _idx_list_struct_equal(x.owner.op.idx_list, node.op.idx_list)


Hmmm, so we should simply start counting from the first indexing, regardless of whether it's Subtensor or AdvancedSubtensor? That would allow us to compare the idx_list directly.

Start both at zero. x[idx] and x[idx].set(y) both have idx_list == [0]

…bility

jaanerik · 2026-02-04T11:10:24Z

Hmmm, so we should simply start counting from the first indexing, regardless of whether it's Subtensor or AdvancedSubtensor? That would allow us to compare the idx_list directly.

Start both at zero. x[idx] and x[idx].set(y) both have idx_list == [0]

Refactored to use ==. Could not respond to or resolve some comments directly, idk why.

Removed many comments, removed most makeslice, slicetype usage, but refactoring xtensor is a bigger change and is still quite verbose. Would like some feedback whether I should touch that code at all before continuing :) @ricardoV94

ricardoV94 · 2026-02-05T12:06:08Z

Let's try not to touch XTensor if we can, to keep this PR moving along. If it's using SliceTypes let it remain for those

ricardoV94 · 2026-02-05T12:09:29Z

pytensor/link/numba/dispatch/subtensor.py

+        # Account for scalar indices before it that remove dimensions
+        for i in range(out_adv_axis_pos):
+            if not isinstance(reconstructed_indices[i], slice):
+                out_adv_axis_pos -= 1


I still don't understand this. This dispatch is just for AdvancedSubtensor right? In that case scalar indices count as advanced indices so there shouldn't be any need to do the -=1. Can you find what test triggers this condition?

ricardoV94 · 2026-02-05T12:12:01Z

pytensor/tensor/random/rewriting/basic.py

-        # Note: For simplicity this also excludes subtensor-related expand_dims (np.newaxis).
-        #  If we wanted to support that we could rewrite it as subtensor + dimshuffle
-        #  and make use of the dimshuffle lift rewrite
-        # TODO: This rewrite is aborting with dummy indexing dimensions which aren't a problem


I think this last TODO is still correct. It was meant for stuff like x[[0],], that's equivalent to x[0], in the sense that it can only ever index one entry and won't lead to duplicates, but looks multidimensional.

ricardoV94 · 2026-02-05T12:13:42Z

pytensor/tensor/rewriting/subtensor.py

+    elif len(shape_parts) == 1:
        shape = shape_parts[0]


Is this a bug fix? If so can we add a regression test?

ricardoV94 · 2026-02-05T12:15:11Z

pytensor/tensor/rewriting/subtensor.py

 @node_rewriter(tracks=[AdvancedSubtensor, AdvancedIncSubtensor])
-def bool_idx_to_nonzero(fgraph, node):
-    """Convert boolean indexing into equivalent vector boolean index, supported by our dispatch
+def ravel_multidimensional_bool_idx(fgraph, node):


This seems to be a conflict not well resolved from main? Adding feature that was perhaps removed/stripped down?

ricardoV94

This is looking nearly done!

I saw you did some changes to xtensor. You mentioned that was growing out of scope so I didn't review. Let me know how you want to proceed there, or if you need guidance.

I flagged some changes that seem to be a conflict from main mis-resolved?

Otherwise it's really shaping up!!!

ricardoV94 · 2026-02-05T12:18:31Z

tests/tensor/test_variable.py

    [
-        (tensor3(), (np.newaxis, slice(None), np.newaxis), ("x", 0, "x", 1, 2)),
-        (cscalar(), (np.newaxis,), ("x",)),
+        (tensor3(), (None, slice(None), None), ("x", 0, "x", 1, 2)),


These tests can stay as np.newaxis to reduce diff, but it's a nit given the scope of the PR

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 added 10 commits January 30, 2026 12:33

Fix rewriting, use existing functions, respect subclasses

75e869d

Fix tests

bd6824c

Implement BaseSubtensor

9328582

Fix rebase

3be09cc

Remove one test that incorrectly extends materialized Ops

7407183

Rename tensor_inputs to index_variables

5bca868

Fix JAX dispatch

cd8f912

Refactor AdvancedSubtensor1 to use idx_list, add comments

c6d80f6

Implement simpler idx_list

e1461c2

Revert unrelated code, remove deprecated code from a test

2cf80f6

jaanerik force-pushed the refactor-advanced-subtensor branch 2 times, most recently from 0003a80 to ff2b297 Compare January 30, 2026 11:23

ricardoV94 reviewed Jan 30, 2026

View reviewed changes

jaanerik added 2 commits January 30, 2026 14:14

Remove symbolic slices

95ef11f

Fix leaky test

dadeedf

jaanerik force-pushed the refactor-advanced-subtensor branch 2 times, most recently from 8ba1d7c to c84885b Compare February 2, 2026 11:26

jaanerik added 2 commits February 3, 2026 16:51

Simplify

43d7e59

Add xtensor hashing and restore is_full_slice for pymc-extras compati…

e127e4b

…bility

jaanerik force-pushed the refactor-advanced-subtensor branch from c84885b to e127e4b Compare February 3, 2026 14:52

Remove comments and useless code

13d025a

jaanerik force-pushed the refactor-advanced-subtensor branch from 0b51466 to 13d025a Compare February 4, 2026 12:48

ricardoV94 reviewed Feb 5, 2026

View reviewed changes

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!

ricardoV94 left a comment

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Choose a reason for hiding this comment

Uh oh!

jaanerik commented Feb 4, 2026

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

ricardoV94 Feb 5, 2026 •

edited

Loading