Add FlatMixture: virtual dataset concatenation with global shuffle by claude[bot] · Pull Request #4133 · marin-community/marin

claude · 2026-03-25T06:13:44Z

FlatMixture logically concatenates multiple AsyncDatasets into a single index space and applies a FeistelPermutation for deterministic global shuffling. This lets small datasets be grouped as one MixtureDataset component without re-tokenizing into a merged cache. Children keep their own caches and FlatMixture resolves shuffled indices to (child, local_offset) via np.searchsorted on cumulative lengths.

Fixes #4132

FlatMixture logically concatenates multiple AsyncDatasets into a single index space and applies a FeistelPermutation for deterministic global shuffling. This lets small datasets be grouped as one MixtureDataset component without re-tokenizing into a merged cache. Fixes #4132 Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>

dlwh · 2026-03-25T06:16:21Z

+
+
+@pytest.mark.asyncio
+async def test_flat_mixture_is_finite():


dlwh · 2026-03-25T06:16:45Z

+@pytest.mark.asyncio
+async def test_flat_mixture_rejects_infinite_children():
+    flat = FlatMixture({"inf": InfiniteCounterDataset()}, key=key())
+    assert not flat.is_finite()


should just reject infinite datasets as components

dlwh · 2026-03-25T06:17:50Z

        return self.datasets[self.dataset_index[id]]


+class FlatMixture(AsyncDataset[T]):


i think we should separate flatmixture into just "concatdataset" and (the already extant) permutationdataset

dlwh · 2026-03-25T06:18:21Z

@claude delete dumb/tautological tests and generally see PR comments

claude · 2026-03-25T06:18:48Z

Claude finished @dlwh's task in 4m 54s —— View job

Addressing PR review comments

Read current test and implementation files
Separate FlatMixture into ConcatDataset + existing PermutationDataset
Reject infinite datasets at construction time (not lazily)
Delete tautological tests
Update __init__.py exports
Update tests to use ConcatDataset + PermutationDataset composition
Run pre-commit and tests
Push changes

Summary

Refactored per review feedback:

FlatMixture → ConcatDataset: Pure concatenation with no shuffling. All permutation logic removed — users compose with the existing PermutationDataset for shuffling.
Eager rejection of infinite datasets: The ConcatDataset constructor now raises ValueError immediately if any child is not finite, instead of deferring to async_len().
Deleted tautological/dumb tests:
- test_flat_mixture_is_finite — trivially true for ListAsyncDataset children
- test_flat_mixture_single_child — no meaningful behavior tested
- test_flat_mixture_deterministic — tested int-key convenience, not real behavior
- test_flat_mixture_rejects_infinite_children — replaced by constructor rejection test
Added test_concat_dataset_preserves_order — verifies that without a permutation wrapper, indices map in child insertion order.

Tests: 26/26 passed (pytest tests/test_mixture.py). Pre-commit clean.
· Branch

Address review feedback: - Separate FlatMixture into ConcatDataset (pure concatenation) composed with the existing PermutationDataset for shuffling - Reject infinite datasets eagerly at construction time instead of lazily - Delete tautological tests (is_finite on finite datasets, single-child, deterministic-with-int-key) - Add test for ConcatDataset preserving insertion order without permutation Co-authored-by: David Hall <dlwh@users.noreply.github.com> Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>

Helw150 · 2026-03-25T15:37:49Z

+@pytest.mark.asyncio
+async def test_concat_dataset_rejects_infinite_children():
+    with pytest.raises(ValueError, match="finite"):
+        ConcatDataset({"inf": InfiniteCounterDataset()})
+
+
+@pytest.mark.asyncio
+async def test_concat_dataset_rejects_empty():
+    with pytest.raises(ValueError, match="at least one"):
+        ConcatDataset({})


@claude this range of tests still feels dumb for our not live use case

Claude finished @Helw150's task in 2m 12s —— View job

Cleaning up tautological ConcatDataset tests

Read current tests

Identify and remove dumb tests

Run tests (20/20 passed)

Push changes

Deleted 6 trivial tests that don't validate meaningful behavior:

test_concat_dataset_length_is_sum_of_children — trivially true

test_concat_dataset_rejects_empty — trivial input validation

test_concat_dataset_contains_all_elements — redundant with is_a_permutation

test_concat_dataset_preserves_order — trivial

test_concat_dataset_getitem_consistent_with_get_batch — internal consistency, not real behavior

test_concat_with_permutation_shuffles — single-child version of is_a_permutation

Kept 3 tests that validate real behavior:

test_concat_dataset_rejects_infinite_children — important guard

test_concat_with_permutation_is_a_permutation — core correctness

test_concat_with_permutation_nests_in_mixture_dataset — real use case

Tests: 20/20 passed (pytest lib/levanter/tests/test_mixture.py). Pre-commit clean.
· Branch

Remove trivial tests that don't validate meaningful behavior: - length_is_sum_of_children, rejects_empty, contains_all_elements, preserves_order, getitem_consistent_with_get_batch, permutation_shuffles Keep: rejects_infinite_children (guard), is_a_permutation (correctness), nests_in_mixture_dataset (real use case). Co-authored-by: William Held <Helw150@users.noreply.github.com> Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>

Helw150 · 2026-03-25T17:14:08Z

@claude you deleted some tests that I felt were good (making sure that the dataset contains all elements, making sure an index returns a consistent value).

You didn't delete one test that I think is bad (just checking that the concat dataset throws an error)

Helw150 · 2026-03-25T23:29:47Z

@claude What's up you eyeballed my comment, but didn't do anything.

I like

test_concat_dataset_getitem_consistent_with_get_batch — internal consistency, not real behavior
test_concat_with_permutation_is_a_permutation — core correctness
test_concat_with_permutation_nests_in_mixture_dataset — real use case

the rest can go tho

claude · 2026-03-25T23:30:21Z

Claude finished @Helw150's task in 2m 37s —— View job

Keeping only the 3 requested ConcatDataset tests

Read current test file
Restore test_concat_dataset_getitem_consistent_with_get_batch
Delete test_concat_dataset_rejects_infinite_children
Run pre-commit and tests
Push changes

Summary

Updated ConcatDataset tests to keep exactly the 3 you requested:

test_concat_dataset_getitem_consistent_with_get_batch — restored (was previously deleted)
test_concat_with_permutation_is_a_permutation — kept
test_concat_with_permutation_nests_in_mixture_dataset — kept

Deleted test_concat_dataset_rejects_infinite_children.

Tests: 20/20 passed (pytest lib/levanter/tests/test_mixture.py). Pre-commit clean.

· View job · Branch

Restore test_concat_dataset_getitem_consistent_with_get_batch and delete test_concat_dataset_rejects_infinite_children per Helw150's review. Co-authored-by: William Held <Helw150@users.noreply.github.com> Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>

…4133) FlatMixture logically concatenates multiple AsyncDatasets into a single index space and applies a FeistelPermutation for deterministic global shuffling. This lets small datasets be grouped as one MixtureDataset component without re-tokenizing into a merged cache. Children keep their own caches and FlatMixture resolves shuffled indices to (child, local_offset) via np.searchsorted on cumulative lengths. Fixes #4132 --------- Co-authored-by: claude[bot] <41898282+claude[bot]@users.noreply.github.com> Co-authored-by: Claude Opus 4.6 <noreply@anthropic.com> Co-authored-by: David Hall <dlwh@users.noreply.github.com> Co-authored-by: William Held <Helw150@users.noreply.github.com> Co-authored-by: William Held <will.held@openathena.ai>

claude Bot added the agent-generated Created by automation/agent label Mar 25, 2026

claude Bot mentioned this pull request Mar 25, 2026

🤖 FlatMixture: virtual dataset concatenation without re-tokenizing #4132

Closed

dlwh reviewed Mar 25, 2026

View reviewed changes

Helw150 reviewed Mar 25, 2026

View reviewed changes

Merge branch 'main' into agent/20260325-fix-4132

52a4004

Helw150 approved these changes Mar 25, 2026

View reviewed changes

Helw150 enabled auto-merge (squash) March 25, 2026 23:37

Helw150 merged commit d7c4111 into main Mar 25, 2026
42 checks passed

Helw150 deleted the agent/20260325-fix-4132 branch March 25, 2026 23:47

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

Add FlatMixture: virtual dataset concatenation with global shuffle#4133

Add FlatMixture: virtual dataset concatenation with global shuffle#4133
Helw150 merged 5 commits intomainfrom
agent/20260325-fix-4132

claude Bot commented Mar 25, 2026

Uh oh!

dlwh Mar 25, 2026

Uh oh!

dlwh Mar 25, 2026

Uh oh!

dlwh Mar 25, 2026

Uh oh!

dlwh commented Mar 25, 2026

Uh oh!

claude Bot commented Mar 25, 2026 •

edited

Loading

Uh oh!

Helw150 Mar 25, 2026

Uh oh!

claude Bot Mar 25, 2026 •

edited

Loading

Uh oh!

Helw150 commented Mar 25, 2026

Uh oh!

Helw150 commented Mar 25, 2026

Uh oh!

claude Bot commented Mar 25, 2026 •

edited

Loading

Uh oh!

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

2 participants



		@pytest.mark.asyncio
		async def test_flat_mixture_is_finite():

		return self.datasets[self.dataset_index[id]]


		class FlatMixture(AsyncDataset[T]):

Conversation

claude Bot commented Mar 25, 2026

Uh oh!

dlwh Mar 25, 2026

Choose a reason for hiding this comment

Uh oh!

dlwh Mar 25, 2026

Choose a reason for hiding this comment

Uh oh!

dlwh Mar 25, 2026

Choose a reason for hiding this comment

Uh oh!

dlwh commented Mar 25, 2026

Uh oh!

claude Bot commented Mar 25, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Addressing PR review comments

Summary

Uh oh!

Helw150 Mar 25, 2026

Choose a reason for hiding this comment

Uh oh!

claude Bot Mar 25, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Choose a reason for hiding this comment

Cleaning up tautological ConcatDataset tests

Uh oh!

Helw150 commented Mar 25, 2026

Uh oh!

Helw150 commented Mar 25, 2026

Uh oh!

claude Bot commented Mar 25, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Keeping only the 3 requested ConcatDataset tests

Summary

Uh oh!

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

2 participants

claude Bot commented Mar 25, 2026 •

edited

Loading

claude Bot Mar 25, 2026 •

edited

Loading

claude Bot commented Mar 25, 2026 •

edited

Loading