Support arbitrary block index functions in blockwise #350
Add this suggestion to a batch that can be applied as a single commit.
This suggestion is invalid because no changes were made to the code.
Suggestions cannot be applied while the pull request is closed.
Suggestions cannot be applied while viewing a subset of changes.
Only one suggestion per line can be applied in a batch.
Add this suggestion to a batch that can be applied as a single commit.
Applying suggestions on deleted lines is not supported.
You must change the existing code in this line in order to create a valid suggestion.
Outdated suggestions cannot be applied.
This suggestion has been applied or marked resolved.
Suggestions cannot be applied from pending reviews.
Suggestions cannot be applied on multi-line comments.
Suggestions cannot be applied while the pull request is queued to merge.
Suggestion cannot be applied right now. Please check back later.
The idea here is to generalize blockwise slightly so you express a dependency of an output array block on input array blocks that you can't express using the index notation.
The example in the test is a simple form of
merge_chunks
where three input blocks are merged into one output block. (This will be useful in improvingreduction
, but that is not included in this PR.)The reason we need this change is so that we can improve fusion optimizations. Currently
merge_chunks
uses the opaquemap_direct
function, which i) assumes that the array being merged has been materialized as a Zarr array, and ii) doesn't give any information about block mapping. By re-writingmerge_chunks
as a blockwise operation we get information about the block mapping that we will be able to use to fuse themerge_chunks
with earlier operations in the DAG.