feat(py): sweep the reassembly axis whose blend window is smallest - #74
Conversation
The window costs `patch[axis] x (the product of the other extents)`, so which axis it slides along decides its size -- and it was axis 0, in the code, in eight places. On 514x1331x1775 with a 320 patch, sweeping axis 2 instead would hold 2.45 GiB where axis 0 holds 8.45: the same volume, 3.45x less resident, because the cut through the volume is smaller. On a 295x259x219 case axis 0 is already the cheapest -- it is a property of each geometry, not a setting. Nothing chooses yet. `sweep_axis` defaults to 0, so every caller gets exactly what it got before; this only moves the axis out of the code and into a value. Two helpers carry it: `_along_sweep` indexes a result over a span on that axis and whole on the others, `_sweep_shape` cuts the spatial shape down to the window. `_advance_to`, the arrival guard, the window sizing and the footprint all read it instead of assuming. A test reassembles the same volume sweeping each of the three axes and checks the footprint is the window on that axis and the extent on the others. Re-hardcoding axis 0 in `_advance_to` fails it on axis 2, so it discriminates. Choosing the axis needs the patch grid emitted with it outermost -- the read and write orders must agree (AGENTS.md) -- and needs measuring: the grid order is also the read order, so a cheaper window can cost read locality. Separate change.
The window slides along one axis, and a patch's arrival finalizes everything behind it on that axis only, so the grid has to be emitted with that axis outermost or the window has nothing to slide on. Both generators now take the axis and route through one helper; the tuples come back in array order, so nothing downstream sees the permutation. Still 0 everywhere. This is the other half of moving the axis out of the code: the accumulator reads it, and now the grid honours it. The order is the contract between the read and the write (AGENTS.md) -- a grid emitted for one axis and reassembled along another hands out regions that are not final, and nothing reports it. So the test asserts the SET of patches is unchanged, not merely the order, and that starts on the swept axis never decrease; axis 0 must reproduce the historical order exactly. Also drops the accumulator list both generators no longer fill.
The window holds min(patch, extent) rows of the swept axis across the whole of every other, so it costs min(patch_d, extent_d) x prod(extent_e != d) -- smallest on the axis the patch divides most. It was axis 0 by construction, which is the cheapest only by accident. Measured on a 256x512x640 volume with a 128 patch and 3 channels: sweep axis 0 (before) 0.47 GiB window 8.16s to read every patch sweep axis 2 (now) 0.19 GiB window 8.03s 2.5x less resident for the same read time. The read does not care: the set of patches is identical, only the order changes, and every chunk is read either way -- which is the opposite of what I expected and the reason this is worth taking. The grid decides and the reassembly asks it (Patch.get_sweep_axis): one source of truth, because a grid emitted for one axis and reassembled along another hands out regions that are not final, and nothing reports it. The choice is tested on five geometries -- the cheapest axis, an already-optimal axis 0, a tie falling back to the first, an axis the patch spans whole (the worst possible sweep), and a free axis, which spans the extent and so never wins. DummyPatch in the streamed-write tests gains the method it doubles.
`Patch.get_sweep_axis` is part of the interface the predictor asks for, so the doubles that stand in for a Patch have to answer it. Four test modules define their own, and only one had been updated -- the other three failed the whole predictor path with an AttributeError. They all return 0: the slices these doubles hand out are cut along axis 0, which is what their assertions assume.
|
Warning Review limit reachedYou’ve reached a temporary PR review limit under our Fair Usage Limits Policy. Next review available in: 8 minutes Your organization has reached its usage spending cap. Adjust your spending cap in the billing tab. How can I continue?After more reviews become available, a review can be triggered using the To avoid repeated limits, reduce automatic review volume by pausing incremental auto-reviews earlier, using label-based review opt-in, excluding WIP or generated PR titles, or requesting reviews manually when the PR is ready. If your team needs uninterrupted high-volume reviews, an organization admin can enable usage-based reviews. How do review limits work?CodeRabbit enforces per-developer PR review limits for each organization. Most developers receive the normal plan review availability. For paid Pro and Pro+ PR reviews, CodeRabbit uses adaptive limits for sustained high-volume activity. When a developer's recent PR review activity reaches the 95th percentile or higher among CodeRabbit users, additional reviews become available more gradually as earlier reviews age out of the rolling window. Please refer docs for additional details. Review details⚙️ Run configurationConfiguration used: Organization UI Review profile: CHILL Plan: Pro Run ID: 📒 Files selected for processing (2)
📝 WalkthroughWalkthroughPatch generation now selects a sweep axis per copy, orders patches accordingly, and passes the axis through predictor-created accumulators. Accumulation and streaming window logic support nonzero axes, with tests covering reconstruction, ordering, selection, and existing test doubles. ChangesConfigurable sweep-axis processing
Estimated code review effort: 3 (Moderate) | ~25 minutes Possibly related PRs
Sequence Diagram(s)sequenceDiagram
participant Patch
participant OutSameAsGroupDataset
participant StreamingAccumulator
Patch->>OutSameAsGroupDataset: provide per-copy sweep_axis
OutSameAsGroupDataset->>StreamingAccumulator: construct accumulator with sweep_axis
Patch->>StreamingAccumulator: provide sweep-first patch slices
StreamingAccumulator-->>OutSameAsGroupDataset: finalize accumulated output
Poem
🚥 Pre-merge checks | ✅ 3 | ❌ 2❌ Failed checks (2 warnings)
✅ Passed checks (3 passed)
✨ Finishing Touches 💡 1📝 Generate docstrings 💡
🧪 Generate unit tests (beta)
Comment |
|
@coderabbitai review |
✅ Action performedReview finished.
|
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
There was a problem hiding this comment.
Actionable comments posted: 1
🤖 Prompt for all review comments with AI agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
Inline comments:
In `@konfai/predictor.py`:
- Line 645: Update the streamed prediction path around StreamingAccumulator and
_consume_slabs so the selected sweep axis remains consistent through sink
writes, buffered writes, and TTA handling; do not apply region as axis 0 when
axes 1 or 2 are selected. If those paths cannot support nonzero axes, route them
through the whole-volume fallback instead. Add a predictor-level streamed test
covering a nonzero selected sweep axis.
🪄 Autofix (Beta)
Fix all unresolved CodeRabbit comments on this PR:
- Push a commit to this branch (recommended)
- Create a new PR with the fixes
ℹ️ Review info
⚙️ Run configuration
Configuration used: Organization UI
Review profile: CHILL
Plan: Pro
Run ID: cf391c60-a72a-4278-9e33-cd3bbb5d49ea
📒 Files selected for processing (7)
konfai/data/patching.pykonfai/predictor.pykonfai/utils/utils.pytests/unit/test_patching.pytests/unit/test_predictor_memory.pytests/unit/test_streamed_tta.pytests/unit/test_streamed_write_dispatcher.py
The streamed consumers (sink target, region pipe, buffer) index their slabs on the first spatial axis, and the aligner needs every augmentation on the same footing. A grid ordered along another axis therefore takes the whole-volume path -- correct if not minimal -- until the consumers carry the axis too.
1 participant
Four commits on top of #70 (now merged): the reassembly window slides along a declared axis instead of an implicit one, the patch grid is emitted with the sweep axis outermost, and the sweep picks the axis whose window is smallest. Every
Patchtest doubles the sweep axis it now declares.Pushed as part of release preparation; CI + CodeRabbit are the gate.
Summary by CodeRabbit
New Features
Bug Fixes
Tests