mixle.data.partition module

Structure-aware partitioning – one place that splits a dataset into chunks while honoring its SampleStructure.

This replaces the inline striding in seq_encode (data[i::C]). For strideable structures (IID / exchangeable / sequential-records) it reproduces that striding exactly, so the existing fast path is bit-identical. For PARTIALLY_EXCHANGEABLE data it strides at the group level – every record of a group lands in the same partition – so a hierarchical model never sees a group split across chunks.

partition_records(records, structure, n)[source]

Split records into n partitions respecting structure.

Strideable structures give records[k::n] (identical to the historical chunking); a partially-exchangeable structure groups by its key and round-robins whole groups across partitions.

Parameters:
Return type:

list[list[Any]]

num_chunks_for(size, num_chunks=1, chunk_size=None)[source]

Resolve the chunk count from an explicit num_chunks or a target chunk_size (as seq_encode does).

Parameters:

  • size (int)

  • num_chunks (int)

  • chunk_size (int | None)

Return type:

int

encode_partitions(records, encoder, structure=EXCHANGEABLE, num_chunks=1, chunk_size=None)[source]

Partition records by structure and encoder.seq_encode each part -> [(count, payload)].

Parameters:

  • records (Sequence[Any])

  • encoder (Any)

  • structure (SampleStructure)

  • num_chunks (int)

  • chunk_size (int | None)

Return type:

list[tuple[int, Any]]