Считайте только токены, действительно вошедшие в мини-пакет

pub fn learner_evidence() -> LearnerEvidence {
    let epoch = build_epoch(SHUFFLE_SEED);
    let replay = build_epoch(SHUFFLE_SEED);
    let changed = build_epoch(SHUFFLE_SEED + 1);
    let batches = epoch
        .batches()
        .iter()
        .enumerate()
        .map(|(index, batch)| batch_evidence(index, batch))
        .collect();

    let ordered_origins = origins(&epoch);
    let mut unique_origins = ordered_origins.clone();
    unique_origins.sort_unstable();
    unique_origins.dedup();
    let expected_origins = [(0, 0), (0, 1), (0, 2), (1, 0), (1, 1)];
    let covered_windows = unique_origins
        .iter()
        .filter(|origin| expected_origins.contains(origin))
        .count();
    let duplicate_windows = ordered_origins.len() - unique_origins.len();
    let padding_ids_added = epoch
        .batches()
        .iter()
        .map(|batch| {
            let expected_stack_len = batch.batch_width() * batch.context_length();
            batch.inputs().len().saturating_sub(expected_stack_len)
                + batch.targets().len().saturating_sub(expected_stack_len)
        })
        .sum();
    let cross_partition_windows = epoch
        .batches()
        .iter()
        .flat_map(MiniBatch::provenance)
        .filter(|origin| origin.partition() != Partition::Train)
        .count();
    LearnerEvidence {
        replay_matches: epoch == replay,
        different_seed_changes_order: origins_in_order(&epoch) != origins_in_order(&changed),
        complete_coverage: unique_origins == expected_origins && duplicate_windows == 0,
        covered_windows,
        expected_windows: expected_origins.len(),
        duplicate_windows,
        padding_ids_added,
        cross_partition_windows,
        epoch,
        batches,
    }
}

fn build_epoch(seed: u64) -> MiniBatchEpoch {
    let documents = [
        BatchDocument::new("train-a", Partition::Train, TRAIN_A).expect("fixture ID is valid"),
        BatchDocument::new("train-b", Partition::Train, TRAIN_B).expect("fixture ID is valid"),
    ];
    let windows = CausalWindowConfig::new(CONTEXT_LENGTH, 1).expect("positive window sizes");
    let batches = MiniBatchConfig::new(REQUESTED_BATCH_SIZE, BatchOrder::Shuffled { seed })
        .expect("positive batch size");
    MiniBatchEpoch::build(Partition::Train, &documents, windows, batches)
        .expect("separate training documents make complete batches")
}

/// Contrasts one-example, three-example, and full-set update widths.
pub fn historical_update_widths(example_count: usize) -> [Vec<usize>; 3] {
    let online = group_widths(example_count, NonZeroUsize::MIN);
    let mini_batch = group_widths(
        example_count,
        NonZeroUsize::new(3).expect("three is positive"),
    );
    let full_capacity = NonZeroUsize::new(example_count).unwrap_or(NonZeroUsize::MIN);
    let full_batch = group_widths(example_count, full_capacity);
    [online, mini_batch, full_batch]
}

fn group_widths(example_count: usize, capacity: NonZeroUsize) -> Vec<usize> {
    let mut widths = Vec::new();
    let mut remaining = example_count;
    while remaining > 0 {
        let width = remaining.min(capacity.get());
        widths.push(width);
        remaining -= width;
    }
    widths
}

/// The stable order used for one materialized epoch.
#[derive(Clone, Copy, Debug, PartialEq, Eq)]
pub enum BatchOrder {
    /// Retain document order, then increasing window start within each document.
    Sequential,
    /// Apply a deterministic Fisher-Yates permutation using the supplied seed.
    Shuffled { seed: u64 },
}

/// A positive requested batch width plus its epoch-order policy.
#[derive(Clone, Copy, Debug, PartialEq, Eq)]
pub struct MiniBatchConfig {
    batch_size: usize,
    order: BatchOrder,
}

impl MiniBatchConfig {
    pub const fn new(batch_size: usize, order: BatchOrder) -> Result<Self, BatchError> {
        if batch_size == 0 {
            return Err(BatchError::ZeroBatchSize);
        }
        Ok(Self { batch_size, order })
    }

    pub const fn batch_size(self) -> usize {
        self.batch_size
    }

    pub const fn order(self) -> BatchOrder {
        self.order
    }
}

/// One separately owned document exposed to the batch builder by reference.
#[derive(Clone, Copy, Debug, PartialEq, Eq)]
pub struct BatchDocument<'a> {
    id: &'a str,
    partition: Partition,
    token_ids: &'a [u32],
}

impl<'a> BatchDocument<'a> {
    pub fn new(
        id: &'a str,
        partition: Partition,
        token_ids: &'a [u32],
    ) -> Result<Self, BatchError> {
        if id.is_empty() {
            return Err(BatchError::EmptyDocumentId);
        }
        Ok(Self {
            id,
            partition,
            token_ids,
        })
    }

    /// Borrows the already-validated provenance and token IDs of one encoded document.
    pub fn from_encoded(document: &'a EncodedDocument) -> Self {
        Self {
            id: document.id(),
            partition: document.partition(),
            token_ids: document.token_ids(),
        }
    }

    pub const fn id(self) -> &'a str {
        self.id
    }

    pub const fn partition(self) -> Partition {
        self.partition
    }

    pub const fn token_ids(self) -> &'a [u32] {
        self.token_ids
    }
}

/// The immutable origin of one complete causal window.
#[derive(Clone, Debug, PartialEq, Eq)]
pub struct WindowProvenance {
    partition: Partition,
    document_index: usize,
    document_id: String,
    start: usize,
}

impl WindowProvenance {
    pub const fn partition(&self) -> Partition {
        self.partition
    }

    pub const fn document_index(&self) -> usize {
        self.document_index
    }

    pub fn document_id(&self) -> &str {
        &self.document_id
    }

    pub const fn start(&self) -> usize {
        self.start
    }
}

#[derive(Clone, Copy, Debug, PartialEq, Eq)]
struct WindowDescriptor {
    document_index: usize,
    start: usize,
}

/// One row-major `[batch, sequence]` stack with no padding rows or token IDs.
#[derive(Clone, Debug, PartialEq, Eq)]
pub struct MiniBatch {
    partition: Partition,
    context_length: usize,
    inputs: Vec<u32>,
    targets: Vec<u32>,
    provenance: Vec<WindowProvenance>,
}

impl MiniBatch {
    pub const fn partition(&self) -> Partition {
        self.partition
    }

    pub const fn context_length(&self) -> usize {
        self.context_length
    }

    pub fn batch_width(&self) -> usize {
        self.provenance.len()
    }

    pub fn shape(&self) -> [usize; 2] {
        [self.batch_width(), self.context_length]
    }

    pub fn token_count(&self) -> usize {
        self.targets.len()
    }

    pub fn inputs(&self) -> &[u32] {
        &self.inputs
    }

    pub fn targets(&self) -> &[u32] {
        &self.targets
    }

    pub fn provenance(&self) -> &[WindowProvenance] {
        &self.provenance
    }

    pub fn input_row(&self, row: usize) -> Option<&[u32]> {
        let start = row.checked_mul(self.context_length)?;
        let end = start.checked_add(self.context_length)?;
        self.inputs.get(start..end)
    }

    pub fn target_row(&self, row: usize) -> Option<&[u32]> {
        let start = row.checked_mul(self.context_length)?;
        let end = start.checked_add(self.context_length)?;
        self.targets.get(start..end)
    }

    /// Averages one checked loss and parameter-gradient vector per target token.
    pub fn average_token_contributions(
        &self,
        contributions: &[TokenContribution],
    ) -> Result<TokenMean, BatchError> {
        let expected = self.token_count();
        if contributions.len() != expected {
            return Err(BatchError::ContributionCountMismatch {
                expected,
                actual: contributions.len(),
            });
        }
        let gradient_width = contributions
            .first()
            .map(TokenContribution::gradient_width)
            .ok_or(BatchError::EmptyAccumulator)?;
        let mut accumulator = TokenMeanAccumulator::new(gradient_width)?;
        for contribution in contributions {
            accumulator.add_token(contribution)?;
        }
        accumulator.finish()
    }
}

/// Every mini-batch in one reproducible traversal of complete windows.
#[derive(Clone, Debug, PartialEq, Eq)]
pub struct MiniBatchEpoch {
    partition: Partition,
    context_length: usize,
    config: MiniBatchConfig,
    window_count: usize,
    shuffle_state_after: Option<u64>,
    batches: Vec<MiniBatch>,
}

impl MiniBatchEpoch {
    pub fn build(
        partition: Partition,
        documents: &[BatchDocument<'_>],
        window_config: CausalWindowConfig,
        config: MiniBatchConfig,
    ) -> Result<Self, BatchError> {
        validate_documents(partition, documents)?;

        let mut window_count = 0_usize;
        for document in documents {
            window_count = window_count
                .checked_add(window_config.window_count(document.token_ids().len()))
                .ok_or(BatchError::WindowCountOverflow)?;
        }

        let mut descriptors = Vec::new();
        descriptors
            .try_reserve_exact(window_count)
            .map_err(|_| BatchError::AllocationFailed {
                elements: window_count,
            })?;
        for (document_index, document) in documents.iter().copied().enumerate() {
            for window in window_config.windows(document.token_ids()) {
                descriptors.push(WindowDescriptor {
                    document_index,
                    start: window.start(),
                });
            }
        }
        debug_assert_eq!(descriptors.len(), window_count);

        let shuffle_state_after = match config.order() {
            BatchOrder::Sequential => None,
            BatchOrder::Shuffled { seed } => {
                let mut rng = SplitMix64::from_seed(seed);
                fisher_yates(&mut descriptors, &mut rng);
                Some(rng.state())
            }
        };

        let batch_count = if window_count == 0 {
            0
        } else {
            (window_count - 1) / config.batch_size() + 1
        };
        let mut batches = Vec::new();
        batches
            .try_reserve_exact(batch_count)
            .map_err(|_| BatchError::AllocationFailed {
                elements: batch_count,
            })?;

        let context_length = window_config.context_length();
        let required_source_tokens = window_config.required_source_tokens();
        let mut descriptors = descriptors.into_iter();
        let mut remaining = window_count;
        while remaining > 0 {
            let width = remaining.min(config.batch_size());
            let token_count = width
                .checked_mul(context_length)
                .ok_or(BatchError::TokenCountOverflow)?;
            let mut inputs = Vec::new();
            let mut targets = Vec::new();
            let mut provenance = Vec::new();
            inputs
                .try_reserve_exact(token_count)
                .map_err(|_| BatchError::AllocationFailed {
                    elements: token_count,
                })?;
            targets
                .try_reserve_exact(token_count)
                .map_err(|_| BatchError::AllocationFailed {
                    elements: token_count,
                })?;
            provenance
                .try_reserve_exact(width)
                .map_err(|_| BatchError::AllocationFailed { elements: width })?;

            for _ in 0..width {
                let descriptor = descriptors
                    .next()
                    .expect("pre-counted descriptor must exist while batching");
                let document = documents[descriptor.document_index];
                let source_end = descriptor.start + required_source_tokens;
                let source = document
                    .token_ids()
                    .get(descriptor.start..source_end)
                    .expect("descriptor must name one complete causal window");

                inputs.extend_from_slice(&source[..context_length]);
                targets.extend_from_slice(&source[1..]);
                provenance.push(WindowProvenance {
                    partition,
                    document_index: descriptor.document_index,
                    document_id: document.id().to_owned(),
                    start: descriptor.start,
                });
            }
            batches.push(MiniBatch {
                partition,
                context_length,
                inputs,
                targets,
                provenance,
            });
            remaining -= width;
        }

        Ok(Self {
            partition,
            context_length,
            config,
            window_count,
            shuffle_state_after,
            batches,
        })
    }

    pub const fn partition(&self) -> Partition {
        self.partition
    }

    pub const fn context_length(&self) -> usize {
        self.context_length
    }

    pub const fn config(&self) -> MiniBatchConfig {
        self.config
    }

    pub const fn window_count(&self) -> usize {
        self.window_count
    }

    pub fn batch_count(&self) -> usize {
        self.batches.len()
    }

    pub const fn shuffle_state_after(&self) -> Option<u64> {
        self.shuffle_state_after
    }

    pub fn batches(&self) -> &[MiniBatch] {
        &self.batches
    }
}

fn validate_documents(
    partition: Partition,
    documents: &[BatchDocument<'_>],
) -> Result<(), BatchError> {
    for (document_index, document) in documents.iter().enumerate() {
        if document.partition() != partition {
            return Err(BatchError::PartitionMismatch {
                document_index,
                expected: partition,
                actual: document.partition(),
            });
        }
        if let Some(first) = documents[..document_index]
            .iter()
            .position(|candidate| candidate.id() == document.id())
        {
            return Err(BatchError::DuplicateDocumentId {
                id: document.id().to_owned(),
                first,
                repeated: document_index,
            });
        }
    }
    Ok(())
}

fn fisher_yates<T>(values: &mut [T], rng: &mut SplitMix64) {
    for upper_index in (1..values.len()).rev() {
        let selected = sample_below(rng, upper_index + 1);
        values.swap(upper_index, selected);
    }
}

fn sample_below(rng: &mut SplitMix64, exclusive_upper: usize) -> usize {
    debug_assert!(exclusive_upper > 0);
    let bound = exclusive_upper as u64;
    let rejection_threshold = bound.wrapping_neg() % bound;
    loop {
        let draw = rng.next_u64();
        if draw >= rejection_threshold {
            return (draw % bound) as usize;
        }
    }
}

/// One target token's scalar loss and parameter-gradient coordinates.
#[derive(Clone, Debug, PartialEq)]
pub struct TokenContribution {
    loss: f64,
    gradient: Vec<f64>,
}

impl TokenContribution {
    pub fn new(loss: f64, gradient: Vec<f64>) -> Result<Self, BatchError> {
        if !loss.is_finite() {
            return Err(BatchError::NonFiniteLoss { value: loss });
        }
        if gradient.is_empty() {
            return Err(BatchError::ZeroGradientWidth);
        }
        if let Some((coordinate, &value)) = gradient
            .iter()
            .enumerate()
            .find(|(_, value)| !value.is_finite())
        {
            return Err(BatchError::NonFiniteGradient { coordinate, value });
        }
        Ok(Self { loss, gradient })
    }

    pub const fn loss(&self) -> f64 {
        self.loss
    }

    pub fn gradient(&self) -> &[f64] {
        &self.gradient
    }

    pub fn gradient_width(&self) -> usize {
        self.gradient.len()
    }
}

/// Raw sums that can be merged before one final division by token count.
#[derive(Clone, Debug, PartialEq)]
pub struct TokenMeanAccumulator {
    loss_sum: f64,
    gradient_sums: Vec<f64>,
    token_count: usize,
}

impl TokenMeanAccumulator {
    pub fn new(gradient_width: usize) -> Result<Self, BatchError> {
        if gradient_width == 0 {
            return Err(BatchError::ZeroGradientWidth);
        }
        let mut gradient_sums = Vec::new();
        gradient_sums
            .try_reserve_exact(gradient_width)
            .map_err(|_| BatchError::AllocationFailed {
                elements: gradient_width,
            })?;
        gradient_sums.resize(gradient_width, 0.0);
        Ok(Self {
            loss_sum: 0.0,
            gradient_sums,
            token_count: 0,
        })
    }

    pub const fn token_count(&self) -> usize {
        self.token_count
    }

    pub const fn loss_sum(&self) -> f64 {
        self.loss_sum
    }

    pub fn gradient_sums(&self) -> &[f64] {
        &self.gradient_sums
    }

    pub fn add_token(&mut self, contribution: &TokenContribution) -> Result<(), BatchError> {
        if contribution.gradient_width() != self.gradient_sums.len() {
            return Err(BatchError::GradientWidthMismatch {
                expected: self.gradient_sums.len(),
                actual: contribution.gradient_width(),
            });
        }
        let next_count = self
            .token_count
            .checked_add(1)
            .ok_or(BatchError::TokenCountOverflow)?;
        let next_loss = self.loss_sum + contribution.loss();
        if !next_loss.is_finite() {
            return Err(BatchError::NonFiniteAccumulation {
                quantity: "loss",
                coordinate: None,
                value: next_loss,
            });
        }
        validate_gradient_sum(&self.gradient_sums, contribution.gradient())?;

        self.loss_sum = next_loss;
        for (sum, &value) in self.gradient_sums.iter_mut().zip(contribution.gradient()) {
            *sum += value;
        }
        self.token_count = next_count;
        Ok(())
    }

    /// Merges raw sums without averaging either side first.
    pub fn merge(&mut self, other: &Self) -> Result<(), BatchError> {
        if other.gradient_sums.len() != self.gradient_sums.len() {
            return Err(BatchError::GradientWidthMismatch {
                expected: self.gradient_sums.len(),
                actual: other.gradient_sums.len(),
            });
        }
        let next_count = self
            .token_count
            .checked_add(other.token_count)
            .ok_or(BatchError::TokenCountOverflow)?;
        let next_loss = self.loss_sum + other.loss_sum;
        if !next_loss.is_finite() {
            return Err(BatchError::NonFiniteAccumulation {
                quantity: "loss",
                coordinate: None,
                value: next_loss,
            });
        }
        validate_gradient_sum(&self.gradient_sums, &other.gradient_sums)?;

        self.loss_sum = next_loss;
        for (sum, &value) in self.gradient_sums.iter_mut().zip(&other.gradient_sums) {
            *sum += value;
        }
        self.token_count = next_count;
        Ok(())
    }

    pub fn finish(self) -> Result<TokenMean, BatchError> {
        if self.token_count == 0 {
            return Err(BatchError::EmptyAccumulator);
        }
        let denominator = self.token_count as f64;
        let mean_loss = self.loss_sum / denominator;
        let mut mean_gradient = self.gradient_sums;
        for value in &mut mean_gradient {
            *value /= denominator;
        }
        Ok(TokenMean {
            token_count: self.token_count,
            mean_loss,
            mean_gradient,
        })
    }
}

fn validate_gradient_sum(left: &[f64], right: &[f64]) -> Result<(), BatchError> {
    debug_assert_eq!(left.len(), right.len());
    for (coordinate, (&left, &right)) in left.iter().zip(right).enumerate() {
        let value = left + right;
        if !value.is_finite() {
            return Err(BatchError::NonFiniteAccumulation {
                quantity: "gradient",
                coordinate: Some(coordinate),
                value,
            });
        }
    }
    Ok(())
}

/// One scalar mean loss and one equally normalized parameter-gradient vector.
#[derive(Clone, Debug, PartialEq)]
pub struct TokenMean {
    token_count: usize,
    mean_loss: f64,
    mean_gradient: Vec<f64>,
}

impl TokenMean {
    pub const fn token_count(&self) -> usize {
        self.token_count
    }

    pub const fn mean_loss(&self) -> f64 {
        self.mean_loss
    }

    pub fn mean_gradient(&self) -> &[f64] {
        &self.mean_gradient
    }
}

fn batch_contributions(batch: &MiniBatch) -> Vec<TokenContribution> {
    batch
        .provenance()
        .iter()
        .flat_map(|origin| {
            (0..batch.context_length()).map(move |token| {
                let numerator = origin.document_index() * 8 + origin.start() * 2 + token + 1;
                let loss = numerator as f64 / 8.0;
                TokenContribution::new(loss, vec![2.0 * loss, 2.0 - loss])
                    .expect("binary-fraction fixture remains finite")
            })
        })
        .collect()
}

fn batch_evidence(index: usize, batch: &MiniBatch) -> BatchEvidence {
    let contributions = batch_contributions(batch);
    let batch_mean = batch
        .average_token_contributions(&contributions)
        .expect("one contribution exists for every target token");
    let direct = accumulate(&contributions);
    let loss_sum = direct.loss_sum();
    let mean = direct.finish().expect("batch has target tokens");
    assert_eq!(mean, batch_mean);

    let split = contributions.len() / 2;
    let mut left = accumulate(&contributions[..split]);
    let right = accumulate(&contributions[split..]);
    left.merge(&right).expect("gradient widths match");
    let accumulated = left.finish().expect("batch has target tokens");

    BatchEvidence {
        index,
        loss_sum,
        losses: contributions.iter().map(TokenContribution::loss).collect(),
        accumulation_matches: accumulated == mean,
        mean,
    }
}

fn accumulate(contributions: &[TokenContribution]) -> TokenMeanAccumulator {
    let mut accumulator = TokenMeanAccumulator::new(2).expect("two fixture coordinates");
    for contribution in contributions {
        accumulator
            .add_token(contribution)
            .expect("finite fixture contribution");
    }
    accumulator
}

fn main() {
    print!("{}", ch21_mini_batches::learner_report());
}

pub fn diagram_trace() -> String {
    let evidence = learner_evidence();
    let epoch = &evidence.epoch;
    let mut lines = vec![format!(
        "META|context={}|capacity={REQUESTED_BATCH_SIZE}|seed={SHUFFLE_SEED}|windows={}|batches={}",
        epoch.context_length(),
        epoch.window_count(),
        epoch.batch_count(),
    )];

    let mut slot = 0;
    for (batch, batch_evidence) in epoch.batches().iter().zip(&evidence.batches) {
        for (row, origin) in batch.provenance().iter().enumerate() {
            let loss_start = row * batch.context_length();
            let loss_end = loss_start + batch.context_length();
            lines.push(format!(
                "WINDOW|slot={slot}|batch={}|row={row}|document={}|document_index={}|start={}|input={}|target={}|losses={}",
                batch_evidence.index,
                origin.document_id(),
                origin.document_index(),
                origin.start(),
                format_ids(batch.input_row(row).expect("row exists")),
                format_ids(batch.target_row(row).expect("row exists")),
                format_values(&batch_evidence.losses[loss_start..loss_end]),
            ));
            slot += 1;
        }
        lines.push(format!(
            "BATCH|index={}|width={}|shape=[{}, {}]|tokens={}|loss_sum={:.6}|mean_loss={:.6}|mean_gradient={}|accumulation={}",
            batch_evidence.index,
            batch.batch_width(),
            batch.batch_width(),
            batch.context_length(),
            batch.token_count(),
            batch_evidence.loss_sum,
            batch_evidence.mean.mean_loss(),
            format_values(batch_evidence.mean.mean_gradient()),
            if batch_evidence.accumulation_matches { "equal" } else { "different" },
        ));
    }

    let final_batch = epoch.batches().last().expect("fixture has batches");
    lines.push(format!(
        "FINAL|width={}|tokens={}|capacity_tokens={}|actual_denominator={}",
        final_batch.batch_width(),
        final_batch.token_count(),
        REQUESTED_BATCH_SIZE * epoch.context_length(),
        final_batch.token_count(),
    ));
    lines.push(format!(
        "PROOF|coverage={}/{}|duplicates={}|padding={}|cross_partition={}|replay={}|different_seed={}|accumulation={}",
        evidence.covered_windows,
        evidence.expected_windows,
        evidence.duplicate_windows,
        evidence.padding_ids_added,
        evidence.cross_partition_windows,
        if evidence.replay_matches { "same" } else { "different" },
        if evidence.different_seed_changes_order { "changed" } else { "same" },
        if evidence.batches.iter().all(|batch| batch.accumulation_matches) {
            "equal"
        } else {
            "different"
        },
    ));
    lines.join("\n") + "\n"
}