Skip to content

feat: bound class sampler - #253

Draft
selmanozleyen wants to merge 3 commits into
scverse:ig/weighted_samplerfrom
selmanozleyen:feat/bound-class-sampler
Draft

selmanozleyen wants to merge 3 commits into
scverse:ig/weighted_samplerfrom
selmanozleyen:feat/bound-class-sampler

Conversation

@selmanozleyen

@selmanozleyen selmanozleyen commented Jul 13, 2026

Copy link
Copy Markdown
Member

Putting it here to share my WIP only. Will fill the body. I tried to split this into nice 3 commits for easier review

update: it is in ai draft state atm.

@codecov

codecov Bot commented Jul 13, 2026

Copy link
Copy Markdown

Codecov Report

❌ Patch coverage is 94.89489% with 17 lines in your changes missing coverage. Please review.
✅ Project coverage is 91.76%. Comparing base (374a41f) to head (ab1843d).

Files with missing lines Patch % Lines
src/annbatch/samplers/_utils.py 82.35% 9 Missing ⚠️
...h/samplers/_class_samplers/_bound_class_sampler.py 92.59% 6 Missing ⚠️
...nnbatch/samplers/_class_samplers/_class_sampler.py 99.10% 1 Missing ⚠️
.../annbatch/samplers/_class_samplers/_rle_manager.py 98.64% 1 Missing ⚠️
Additional details and impacted files
@@            Coverage Diff             @@
##             main     #253      +/-   ##
==========================================
+ Coverage   91.67%   91.76%   +0.09%     
==========================================
  Files          15       19       +4     
  Lines        1453     1652     +199     
==========================================
+ Hits         1332     1516     +184     
- Misses        121      136      +15     
Files with missing lines Coverage Δ
src/annbatch/samplers/__init__.py 100.00% <100.00%> (ø)
src/annbatch/samplers/_class_samplers/__init__.py 100.00% <100.00%> (ø)
...amplers/_class_samplers/_weighted_class_sampler.py 100.00% <100.00%> (ø)
...nnbatch/samplers/_class_samplers/_class_sampler.py 99.10% <99.10%> (ø)
.../annbatch/samplers/_class_samplers/_rle_manager.py 98.64% <98.64%> (ø)
...h/samplers/_class_samplers/_bound_class_sampler.py 92.59% <92.59%> (ø)
src/annbatch/samplers/_utils.py 88.88% <82.35%> (-8.55%) ⬇️

... and 1 file with indirect coverage changes

🚀 New features to boost your workflow:
  • ❄️ Test Analytics: Detect flaky tests, report on failures, and find test suite problems.

A class's run was picked uniformly among its runs, then a start uniformly inside
it, so every run contributed one equally likely position regardless of size.
With a 400-row run beside a 40-row one, each got half the chunks and the short
run's rows were drawn 10x as often per row. Class marginals were unaffected, so
no share test noticed.

The ClassSampler docstring still describes the intended scheme: "a uniform
chunk-start within c is drawn per chunk (a prefix-sum lookup maps it to the
absolute slice in O(log n_runs))". The sampler worked that way until ec49f9b,
where the prefix sum changed from cumulative valid positions to cumulative run
counts, leaving searchsorted nothing to resolve against. It was replaced by a
direct index plus a second draw.

Each run now records how many chunk starts it holds and how many precede it in
the table, and one uniform draw per class resolves through those to a run and an
offset inside it. Two rng.integers calls become one, so a given seed yields
different slices than before.

Both class samplers draw through RLEManager.slices_from_classes, so this fixes
WeightedClassSampler too, and the new tests run over both.
@selmanozleyen
selmanozleyen force-pushed the feat/bound-class-sampler branch from be22c7a to ab1843d Compare September 18, 2026 23:54
@selmanozleyen
selmanozleyen changed the base branch from main to ig/weighted_sampler September 19, 2026 00:32
@selmanozleyen
selmanozleyen force-pushed the feat/bound-class-sampler branch 3 times, most recently from 89da394 to 11ef92c Compare September 19, 2026 00:55
@selmanozleyen
selmanozleyen force-pushed the feat/bound-class-sampler branch 6 times, most recently from 3c94033 to edbc769 Compare September 19, 2026 16:04
Replays another class sampler's per-batch class schedule against a second
annotation column, so one pass reads two annotations in step with each other.
It takes a ClassSampler or another BoundClassSampler -- so these chain -- and
rejects WeightedClassSampler, whose batches mix classes by design.

The shared machinery moves into two private bases in _class_sampler.py:
_RunClassSampler holds the constructor, the remainder handling and the window
loop, and _ScheduledClassSampler adds the per-batch schedule that makes a
sampler replayable. WeightedClassSampler drops its near-duplicate
_iter_requests and overrides only the two policy hooks, which also removes the
Liskov problem of it subclassing ClassSampler while batching differently.

BoundClassSampler keeps a private copy of its inner sampler, so nothing the
caller does to their own instance afterwards can desynchronise the schedule;
its rng setter forwards to that copy, so re-assignment -- as DistributedSampler
does per rank -- still reaches it.
DistributedSampler stored the sampler it wrapped by reference and then sharded
its mask on every n_batches/validate/_sample call and re-spawned its rng once.
Wrapping therefore left the caller's own sampler permanently rank-sharded, so
reusing it for eval read 1/world_size of the data. Shard a copy instead.

The class samplers reused one row-id buffer for a whole pass and reshuffled it
in place. split_given_size is np.split, which returns views, so every window
handed out the same arrays: list(sampler.sample(n)) showed the last window's
permutation throughout, and a caller writing into one window's split corrupted
the next. Allocate per window.

A mask assigned part-way through a pass was reported by .mask but never read
from, because a pass fixes all of its slices when it starts. Refuse it.

class_weights given as a pandas Series was converted positionally, discarding
the index, so weights silently landed on the wrong classes. Reject it and say
how to align it.

All four are pre-existing in main and reproduce identically on every branch.
@selmanozleyen selmanozleyen added the skip-gpu-ci Whether gpu ci should be skipped label Sep 19, 2026

This branch has not been deployed

No deployments
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

skip-gpu-ci Whether gpu ci should be skipped

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant