feat: Use array format for index-set sequences

[nardi]

Summary

This PR adds a new format for storing index-set sequences during detection, which previously were stored as list[set[int]]. These are collections of sets, that are linked to a certain value involved in the traced computation. There is one set per element of the value, and each set contains the indices of the elements in the input that it depends on.

Storing these as list[set[int]] can become quite slow when large arrays with millions of elements are involved. In this case, the representation is also not very sparse: though there may only be few dependencies, even if there are no dependencies we still have to create millions of empty sets. This can become a memory and performance bottleneck.

The new format does two things:

1. It stores the index-set sequences in a CSR-like format, where all int indices are stored in a single array, and an array of offsets indicates where each set starts and ends. This is a format with very efficient storage and constant-time lookup of a set at a specific index, similar to the previous list.
2. Since the CSR format is not appropriate for many write-operations (e.g. insertion at a specific index), a builder pattern has been applied to collect all the indices inserted into an empty index-set sequence. In practice, this pattern is often already applied, since index set sequences are built once and then only read in further propagation rules.
3. As an extension to the list[set[int]]-like API, the index-set sequences can also be indexed in a vectorized manner (as iss[arr]) and the builder can perform vectorized union operations (as builder[arr1] |= arr2). These also integrate with each other (builder[out_arr] |= iss[in_arr]), which is a common pattern since input dependencies are often propagated as-is.

Further work would be to modify propagation rules to build the index-set sequences in a vectorized manner, supported by the new vectorized API.

There are no public API or otherwise user-facing changes in this PR.

Related issue

Closes #174.

Checklist

• This PR addresses an issue the maintainers have previously agreed on (linked above).
• I understand all code changes in this PR and am able to walk reviewers through them.
• I will not use AI to write answers during the code review (with the exception of spell-checking and translation).
• The code changes match the existing code style, and non-public names are underscore-prefixed.
• Lint, format, and type checks pass (uv run prek run --all-files).
• Tests cover the change and pass (uv run pytest).
• Docs and CHANGELOG.md are updated for user-facing changes.

[adrhill]

I'll take a look at this once I find some free time. Some observations from a quick skim:

Introducing index set mutation

I might be mistaken, but it looks like this PR makes two simultaneous, but orthogonal changes:

• (a) Mutating sets instead of allocating new ones (i.e., using .update instead of |= in all handler files)
• (b) Introducing new set types in _common.py

I'm somewhat surprised that (a) didn't break tests, as mutation should be less conservative than allocating sets. Maybe this is a perk of the Jaxpr language being functional. Maybe our test set doesn't yet cover such edge cases. Maybe I also just didn't skim the code rigorously enough and (a) is somehow enabled by (b).

If this change is in orthogonal to (b), its performance gains should be benchmarked separately from the set types.

New set types

It stores the index-set sequences in a CSR-like format, where all int indices are stored in a single array, and an array of offsets indicates where each set starts and ends.

Could you explain this format to me on a simple example?

You introduce a large amount of new types and abstractions: e.g, IndexSetArray, IndexSetOffsetArrays, IndexSetView, IndexSetSequence, IndexSetSequenceBuilder, IndexSetSequenceBuilderIndexer.
Why do we need so many different new types to replace set[int]? Could you go through the types and explain their purpose?

Side note for @gdalle: you will be delighted to see your beloved DuplicateVector from SCT return in form of IndexSetArray here. ;)

[codecov-commenter]

Codecov Report

❌ Patch coverage is 94.82072% with 13 lines in your changes missing coverage. Please review.
✅ Project coverage is 93.86%. Comparing base (4c9eb35) to head (bcf9ab6).

Files with missing lines	Patch %	Lines
src/asdex/detection/_interpret/_common.py	94.17%	12 Missing ⚠️
src/asdex/detection/_interpret/_elementwise.py	66.66%	1 Missing ⚠️

Additional details and impacted files

@@            Coverage Diff             @@
##             main     #175      +/-   ##
==========================================
+ Coverage   93.84%   93.86%   +0.01%     
==========================================
  Files          61       61              
  Lines        3999     4191     +192     
==========================================
+ Hits         3753     3934     +181     
- Misses        246      257      +11     

☔ View full report in Codecov by Harness.
📢 Have feedback on the report? Share it here.

🚀 New features to boost your workflow:

• ❄️ Test Analytics: Detect flaky tests, report on failures, and find test suite problems.

[nardi]

I'll take a look at this once I find some free time. Some observations from a quick skim:

Introducing index set mutation

I might be mistaken, but it looks like this PR makes two simultaneous, but orthogonal changes:

* (a) Mutating sets instead of allocating new ones (i.e., using `.update` instead of `|=` in all handler files)

* (b) Introducing new set types in `_common.py`

I'm somewhat surprised that (a) didn't break tests, as mutation should be less conservative than allocating sets. Maybe this is a perk of the Jaxpr language being functional. Maybe our test set doesn't yet cover such edge cases. Maybe I also just didn't skim the code rigorously enough and (a) is somehow enabled by (b).

I think (a) is indeed an orthogonal change, that was necessitated by the new objects interacting with regular sets. I think perhaps now it is not needed anymore. The non-mutating alternative would be a = a.union(b).

I think it is indeed just a coincidence that it doesn't break anything, but then again I don't see many situations in which you would use |= but then still have a reference to the previous set somewhere that you need to keep as-is. So I would think if possible update would be preferred to limit the amount of copies.

If this change is in orthogonal to (b), its performance gains should be benchmarked separately from the set types.

Of course, I can split it off into a separate commit, or if it's not needed anymore I can also remove the update changes entirely.

For performance testing, should I include this type of "large constant" calculation test? Or do you have a different suggestion? And how/when should it be run? Every pytest run seems excessive.

New set types

It stores the index-set sequences in a CSR-like format, where all int indices are stored in a single array, and an array of offsets indicates where each set starts and ends.

Could you explain this format to me on a simple example?

You introduce a large amount of new types and abstractions: e.g, IndexSetArray, IndexSetOffsetArrays, IndexSetView, IndexSetSequence, IndexSetSequenceBuilder, IndexSetSequenceBuilderIndexer. Why do we need so many different new types to replace set[int]? Could you go through the types and explain their purpose?

To clarify, set[int] is not being replaced, list[set[int]] is. I think that is an important distinction, if you have only a single set[int] this format doesn't provide any benefit. But of course, when calculating Jacobians over large vector functions you usually don't.

For the separate types:

1. IndexSetArray is the COO-like format used when building the index-set sequence. This is just a numpy array with a special dtype. It is used during the build stage because adding indices and performing bulk operation is efficient, but then it is converted because indexing into it is slow.
2. IndexSetSequenceBuilder is a nice interface to support creating these IndexSetArray objects from a variety of list[set[int]]-like objects. It mainly exposes the builder[i] |= j API that existing code was already using. It also batches as much data conversion operations as possible together in the build step, to avoid performing them during a hot loop.
3. IndexSetSequenceBuilderIndexer just exists to allow for the builder[i] |= j API, which is equivalent to builder[i] = builder[i].__ior__(j). So builder[i] has to return something, and that is this object, which basically just forwards to builder.
4. IndexSetOffsetArrays is the CSR-like format that is used to store the index-set sequence, which can be indexed efficiently but not modified efficiently. It is just an alias to indicate intent for a tuple of numpy arrays.
5. IndexSetSequence is a nice list[set[int]]-like interface to interact with IndexSetOffsetArrays, mainly providing indexing.
6. IndexSetView is the object returned when indexing an IndexSetSequence. To avoid a copy it stores a view of the numpy array, but also enables some interop with regular sets.

Is that clearer? It's a lot of types but I don't really see a way in which removing any of them or folding them into one would make the code simpler, it would only create objects with mixed responsibilities (IMO).

[nardi]

I've removed the |= to update changes, those make more sense to consider when making per-rule performance improvements. I think the changes can just be squashed in, but left them as a fixup commit for now just in case.

[adrhill]

Thanks for the explanations!

1. IndexSetArray is the COO-like format

5. IndexSetOffsetArrays is the CSR-like format

Why do we need both COO and CSR? Where is which one used?

2. IndexSetSequenceBuilder is a nice interface

3. IndexSetSequenceBuilderIndexer just exists to allow for the builder[i] |= j API

5. IndexSetSequence is a nice list[set[int]]-like interface

If 3 out of 6 types (IndexSetSequenceBuilder, IndexSetSequenceBuilderIndexer, IndexSetSequence) are just needed for the internal, non-user facing interface, they should be removed. I'm worried more coupled abstraction (i.e., more wrapper types) will make the code harder to understand and maintain.

6. IndexSetView is the object returned when indexing an IndexSetSequence. To avoid a copy it stores a view of the numpy array, but also enables some interop with regular sets.

Shouldn't all index sets be views then? When are views used over allocating new sets and vice-versa?

If this is just used in intermediate computations within handlers, we might want to make this a class method of IndexSetSequence, since it doesn't own any memory.

I think the changes can just be squashed in

As I've mentioned in #174, I have large concurrent branch of work I did with Fable two weeks ago, so this won't be a simple squash merge. I just made it public as a draft PR in #177 in case you are curious (it's a very large refactor and I haven't had time to review it yet either). Since I'm a strapped for time (I need to write up my PhD thesis this summer), I'll prioritize reviewing and merging #176 first.

This PR is also missing benchmarks demonstrating improvements on both small and large problems (~100k-1M inputs). I worry IndexSetArray might run out of memory on large problems. As I've mentioned in #174:

The issue with selecting a set type is that its performance is highly dependent on the specific problem at hand, mostly the input dimensionality and sparsity. As you've said yourself, benchmarking results are problem-specific. The only clean solution I see is to make the set types selectable by users, like we do in our Julia code.

[Copilot]

Pull request overview

This PR replaces the interpreter’s internal “index-set sequence” representation (list[set[int]]) with a compact, array-backed CSR-like format (IndexSetSequence) plus a builder (IndexSetSequenceBuilder) to reduce memory/time overhead when tracking sparsity for very large arrays.

Changes:

• Introduces IndexSetSequence / IndexSetView / IndexSetSequenceBuilder and updates StateIndices to normalize assigned values into IndexSetSequence.
• Updates multiple propagation rules to use views/builders (|=) instead of mutating or constructing per-element Python sets/lists.
• Adds a dedicated unit test suite for the new containers and adjusts existing tests for the new return types.

Reviewed changes

Copilot reviewed 19 out of 19 changed files in this pull request and generated 4 comments.

Show a summary per file

File	Description
tests/_interpret/test_internals.py	Updates assertions to accommodate IndexSetSequence element views.
tests/_interpret/test_index_set_sequence.py	Adds comprehensive unit tests for IndexSetSequence, views, builder semantics, and StateIndices normalization.
src/asdex/detection/_interpret/CLAUDE.md	Updates internal interpreter docs to describe the new containers and aliasing/immutability rules.
src/asdex/detection/_interpret/_while.py	Ensures while-loop carry is copied into mutable sets before fixed-point mutation; updates typing/contracts.
src/asdex/detection/_interpret/_stack.py	Switches pooled indices container to store IndexSetView rather than mutable sets.
src/asdex/detection/_interpret/_sort.py	Replaces per-element assignment with builder `
src/asdex/detection/_interpret/_scatter.py	Adjusts scatter handler signatures/outputs to work with IndexSetView sequences.
src/asdex/detection/_interpret/_scan.py	Threads carry as sequences of set-likes and concatenates per-step outputs using views.
src/asdex/detection/_interpret/_reduce.py	Uses IndexSetSequenceBuilder for reduction accumulation via `
src/asdex/detection/_interpret/_pad.py	Adjusts output typing to allow returning set-like views/sets without assuming mutability.
src/asdex/detection/_interpret/_gather.py	Updates gather enumeration callback types to return view-based sequences.
src/asdex/detection/_interpret/_elementwise.py	Returns views directly when one derivative is globally zero; avoids unnecessary copies.
src/asdex/detection/_interpret/_dynamic_slice.py	Switches dynamic slice/update to use view replacement rather than deep-copying/mutating sets.
src/asdex/detection/_interpret/_dot_general.py	Uses builder for output accumulation instead of pre-allocating a list of sets.
src/asdex/detection/_interpret/_cond.py	Uses builder-based copying/merging for branch output unions.
src/asdex/detection/_interpret/_concatenate.py	Switches pooled indices container to store IndexSetView.
src/asdex/detection/_interpret/_common.py	Implements core array formats, view/sequence/builder types, updates StateIndices, and adapts common helpers to set-like sequences.
src/asdex/detection/_interpret/init.py	Updates _prop_jaxpr signature/return types and internal state initialization to use StateIndices.
src/asdex/detection/_api.py	Updates input seeding to build IndexSetSequence efficiently (identity/empty builders).

---

💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.

[nardi]

Why do we need both COO and CSR? Where is which one used?

Because COO is efficient when building the sets, and CSR when reading them later. I feel like I am repeating what I said before, is there something still unclear with regards to these structures?

If 3 out of 6 types (IndexSetSequenceBuilder, IndexSetSequenceBuilderIndexer, IndexSetSequence) are just needed for the internal, non-user facing interface, they should be removed. I'm worried more coupled abstraction (i.e., more wrapper types) will make the code harder to understand and maintain.

Okay, so you'd rather I update all set-building code to work with a new interface? I thought it would be preferable to keep the changes as localized as possible, but that works too :)

6. IndexSetView is the object returned when indexing an IndexSetSequence. To avoid a copy it stores a view of the numpy array, but also enables some interop with regular sets.

Shouldn't all index sets be views then? When are views used over allocating new sets and vice-versa?

When reading yes, you are almost always using a view. Building a new set-sequence is the only situation in which you need to take a number of these views and recombine them in a different manner.

If this is just used in intermediate computations within handlers, we might want to make this a class method of IndexSetSequence, since it doesn't own any memory.

What are you referring to here?

I think the changes can just be squashed in

As I've mentioned in #174, I have large concurrent branch of work I did with Fable two weeks ago, so this won't be a simple squash merge. I just made it public as a draft PR in #177 in case you are curious (it's a very large refactor and I haven't had time to review it yet either). Since I'm a strapped for time (I need to write up my PhD thesis this summer), I'll prioritize reviewing and merging #176 first.

I meant, my new changes after your first comments can be squashed into my initial changes, since they make the scope of the MR smaller :)

This PR is also missing benchmarks demonstrating improvements on both small and large problems (~100k-1M inputs). I worry IndexSetArray might run out of memory on large problems. As I've mentioned in #174:

The issue with selecting a set type is that its performance is highly dependent on the specific problem at hand, mostly the input dimensionality and sparsity. As you've said yourself, benchmarking results are problem-specific. The only clean solution I see is to make the set types selectable by users, like we do in our Julia code.

Sure, performance benchmarks would be good to add. I asked before if you have any opinion on what would be a representative benchmark, otherwise I can just add this "large constant" test.

However with regard to your set-type comment, I don't think there would be any situation where IndexSetArray would run out of memory but the equivalent list[set[int]] would not, since it's much more lightweight (e.g. the overhead per set object is larger than one int32). Do you foresee any problems there?