Сохраните состояние декодера и точно повторите одно заданное обновление

#[derive(Clone, Debug, PartialEq, Eq)]
pub struct HistoricalCheckpointContrast {
    pub isolated_parameter_bytes: Vec<u8>,
    pub isolated_parameter_tensors: usize,
    pub isolated_parameter_scalars: usize,
    pub checkpoint_records: usize,
    pub tokenizer_records: usize,
    pub optimizer_moment_records: usize,
    pub checkpoint_file_bytes: usize,
    pub checkpoint_sampling_rng_state: u64,
}

/// Contrasts model-derived value bytes with the complete local LLM checkpoint.
pub fn historical_checkpoint_contrast(
    checkpoint: &Checkpoint,
    encoded: &llm_from_scratch::checkpoint::EncodedCheckpoint,
) -> HistoricalCheckpointContrast {
    let isolated_parameter_bytes = checkpoint
        .model_state()
        .bit_pattern()
        .into_iter()
        .flat_map(u64::to_le_bytes)
        .collect();
    let tokenizer_records = encoded
        .tensors()
        .iter()
        .filter(|descriptor| {
            matches!(
                descriptor.role(),
                CheckpointTensorRole::LiteralToken | CheckpointTensorRole::BpePairs
            )
        })
        .count();
    let optimizer_moment_records = encoded
        .tensors()
        .iter()
        .filter(|descriptor| {
            matches!(
                descriptor.role(),
                CheckpointTensorRole::OptimizerFirstMoment
                    | CheckpointTensorRole::OptimizerSecondMoment
            )
        })
        .count();

    HistoricalCheckpointContrast {
        isolated_parameter_bytes,
        isolated_parameter_tensors: checkpoint.model_state().parameter_names().len(),
        isolated_parameter_scalars: checkpoint.model_state().scalar_count(),
        checkpoint_records: encoded.tensors().len(),
        tokenizer_records,
        optimizer_moment_records,
        checkpoint_file_bytes: encoded.bytes().len(),
        checkpoint_sampling_rng_state: checkpoint.sampling_rng_state(),
    }
}

/// One validated name-keyed moment pair prepared for persistence.
#[derive(Clone, Debug, PartialEq)]
pub struct AdamWStateEntry {
    name: String,
    moments: AdamWMomentState,
}

impl AdamWStateEntry {
    pub fn new(
        name: impl Into<String>,
        shape: Vec<usize>,
        first_moment: Vec<f64>,
        second_moment: Vec<f64>,
    ) -> Result<Self, AdamWStateError> {
        let name = name.into();
        if name.is_empty() {
            return Err(AdamWStateError::EmptyParameterName);
        }
        let elements = shape.iter().try_fold(1_usize, |product, &dimension| {
            product.checked_mul(dimension)
        });
        let Some(elements) = elements else {
            return Err(AdamWStateError::ShapeProductOverflow { name });
        };
        if first_moment.len() != elements || second_moment.len() != elements {
            return Err(AdamWStateError::MomentLengthMismatch {
                name,
                shape,
                expected: elements,
                first: first_moment.len(),
                second: second_moment.len(),
            });
        }
        for (kind, values) in [
            ("first", first_moment.as_slice()),
            ("second", second_moment.as_slice()),
        ] {
            if let Some((index, &value)) = values
                .iter()
                .enumerate()
                .find(|(_, value)| !value.is_finite())
            {
                return Err(AdamWStateError::NonFiniteMoment {
                    name,
                    kind,
                    index,
                    value,
                });
            }
        }
        if let Some((index, &value)) = second_moment
            .iter()
            .enumerate()
            .find(|(_, value)| **value < 0.0)
        {
            return Err(AdamWStateError::NegativeSecondMoment { name, index, value });
        }
        Ok(Self {
            name,
            moments: AdamWMomentState {
                shape,
                first: first_moment,
                second: second_moment,
            },
        })
    }

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

    pub const fn moments(&self) -> &AdamWMomentState {
        &self.moments
    }
}

/// A complete graph-free AdamW continuation snapshot.
#[derive(Clone, Debug, PartialEq)]
pub struct AdamWState {
    config: AdamWConfig,
    groups: Option<AdamWParameterGroups>,
    step: u64,
    beta1_power: f64,
    beta2_power: f64,
    states: BTreeMap<String, AdamWMomentState>,
}

impl AdamWState {
    pub fn new(
        config: AdamWConfig,
        groups: Option<AdamWParameterGroups>,
        step: u64,
        beta1_power: f64,
        beta2_power: f64,
        entries: Vec<AdamWStateEntry>,
    ) -> Result<Self, AdamWStateError> {
        validate_beta_power("beta1", beta1_power, step)?;
        validate_beta_power("beta2", beta2_power, step)?;
        if (step == 0) != entries.is_empty() {
            return Err(AdamWStateError::StatePresence {
                step,
                entries: entries.len(),
            });
        }
        let mut states = BTreeMap::new();
        for entry in entries {
            if states.insert(entry.name.clone(), entry.moments).is_some() {
                return Err(AdamWStateError::DuplicateParameterName { name: entry.name });
            }
        }
        if let Some(groups) = &groups {
            let expected = groups.parameter_names();
            let actual = states.keys().cloned().collect::<Vec<_>>();
            if step > 0 && expected != actual {
                return Err(AdamWStateError::ParameterGroupsMismatch { expected, actual });
            }
        }
        Ok(Self {
            config,
            groups,
            step,
            beta1_power,
            beta2_power,
            states,
        })
    }

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

    pub const fn parameter_groups(&self) -> Option<&AdamWParameterGroups> {
        self.groups.as_ref()
    }

    pub const fn step_count(&self) -> u64 {
        self.step
    }

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

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

    pub fn parameter_names(&self) -> impl ExactSizeIterator<Item = &str> {
        self.states.keys().map(String::as_str)
    }

    pub fn state(&self, name: &str) -> Option<&AdamWMomentState> {
        self.states.get(name)
    }
}

fn validate_beta_power(name: &'static str, value: f64, step: u64) -> Result<(), AdamWStateError> {
    let valid = if step == 0 {
        value.to_bits() == 1.0_f64.to_bits()
    } else {
        value.is_finite() && (0.0..1.0).contains(&value)
    };
    if valid {
        Ok(())
    } else {
        Err(AdamWStateError::InvalidBetaPower { name, value, step })
    }
}

/// A malformed optimizer snapshot rejected before an AdamW instance is built.
#[derive(Clone, Debug, PartialEq)]
pub enum AdamWStateError {
    InvalidBetaPower {
        name: &'static str,
        value: f64,
        step: u64,
    },
    StatePresence {
        step: u64,
        entries: usize,
    },
    EmptyParameterName,
    DuplicateParameterName {
        name: String,
    },
    ShapeProductOverflow {
        name: String,
    },
    MomentLengthMismatch {
        name: String,
        shape: Vec<usize>,
        expected: usize,
        first: usize,
        second: usize,
    },
    NonFiniteMoment {
        name: String,
        kind: &'static str,
        index: usize,
        value: f64,
    },
    NegativeSecondMoment {
        name: String,
        index: usize,
        value: f64,
    },
    ParameterGroupsMismatch {
        expected: Vec<String>,
        actual: Vec<String>,
    },
}

impl fmt::Display for AdamWStateError {
    fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result {
        match self {
            Self::InvalidBetaPower { name, value, step } => write!(
                formatter,
                "{name} power must be exactly 1 at step zero or finite in [0,1) later, got {value} at step {step}"
            ),
            Self::StatePresence { step, entries } => write!(
                formatter,
                "AdamW step {step} is incompatible with {entries} persisted moment entries"
            ),
            Self::EmptyParameterName => {
                formatter.write_str("an AdamW persistence entry has an empty parameter name")
            }
            Self::DuplicateParameterName { name } => write!(
                formatter,
                "AdamW persistence parameter name {name:?} appears more than once"
            ),
            Self::ShapeProductOverflow { name } => write!(
                formatter,
                "AdamW persistence shape for parameter {name:?} overflows usize"
            ),
            Self::MomentLengthMismatch {
                name,
                shape,
                expected,
                first,
                second,
            } => write!(
                formatter,
                "AdamW persistence parameter {name:?} with shape {shape:?} needs {expected} moments, got {first} first and {second} second"
            ),
            Self::NonFiniteMoment {
                name,
                kind,
                index,
                value,
            } => write!(
                formatter,
                "AdamW persistence parameter {name:?} has non-finite {kind} moment at flat index {index}: {value}"
            ),
            Self::NegativeSecondMoment { name, index, value } => write!(
                formatter,
                "AdamW persistence parameter {name:?} has negative second moment at flat index {index}: {value}"
            ),
            Self::ParameterGroupsMismatch { expected, actual } => write!(
                formatter,
                "AdamW persistence groups name {expected:?}, but moments name {actual:?}"
            ),
        }
    }
}

impl Error for AdamWStateError {}

    /// Copies one trainer-issued selected model/AdamW/step bundle.
    ///
    /// The sealed bundle prevents callers from attaching a free step label or an
    /// independently chosen optimizer to the selected model state.
    pub fn from_snapshot(
        tokenizer: CheckpointTokenizer,
        selected: &SelectedTrainingState,
        sampling_rng_state: u64,
    ) -> Result<Self, CheckpointError> {
        let selected_step =
            u64::try_from(selected.step()).map_err(|_| CheckpointError::SizeOverflow {
                context: "selected training step",
            })?;
        Self::from_owned_parts(
            tokenizer,
            selected.model_state().independent_snapshot(),
            selected.optimizer_state().clone(),
            selected_step,
            sampling_rng_state,
        )
    }

    fn from_owned_parts(
        tokenizer: CheckpointTokenizer,
        model_state: DecoderModelState,
        optimizer_state: AdamWState,
        selected_step: u64,
        sampling_rng_state: u64,
    ) -> Result<Self, CheckpointError> {
        let checkpoint = Self {
            tokenizer,
            model_state,
            optimizer_state,
            selected_step,
            sampling_rng_state,
        };
        checkpoint.validate_parts()?;
        Ok(checkpoint)
    }

    pub const fn tokenizer(&self) -> &CheckpointTokenizer {
        &self.tokenizer
    }

    pub const fn model_state(&self) -> &DecoderModelState {
        &self.model_state
    }

    pub const fn optimizer_state(&self) -> &AdamWState {
        &self.optimizer_state
    }

    pub const fn selected_step(&self) -> u64 {
        self.selected_step
    }

    pub const fn sampling_rng_state(&self) -> u64 {
        self.sampling_rng_state
    }

    /// Compatibility accessor for the version-1 sampling RNG state.
    pub const fn rng_state(&self) -> u64 {
        self.sampling_rng_state()
    }

    /// Rebuilds an independent decoder while retaining the checkpoint.
    pub fn restore_independent_model(&self) -> Result<DecoderModel, CheckpointError> {
        self.model_state
            .restore_independent_model()
            .map_err(Into::into)
    }

    /// Consumes the checkpoint and moves its model buffers into one decoder.
    pub fn into_model(self) -> Result<DecoderModel, CheckpointError> {
        self.model_state.into_model().map_err(Into::into)
    }

    pub fn restore_optimizer(&self) -> AdamW {
        AdamW::from_persistence_state(&self.optimizer_state)
    }

#[derive(Clone, Copy, Debug)]
enum TensorPayload<'a> {
    Bytes(&'a [u8]),
    BpePairs(&'a [TokenPair]),
    Float64(&'a [f64]),
}

impl TensorPayload<'_> {
    const fn dtype(self) -> CheckpointDType {
        match self {
            Self::Bytes(_) => CheckpointDType::U8,
            Self::BpePairs(_) => CheckpointDType::U32,
            Self::Float64(_) => CheckpointDType::F64,
        }
    }

    fn byte_len(self) -> Result<usize, CheckpointError> {
        match self {
            Self::Bytes(values) => Ok(values.len()),
            Self::BpePairs(pairs) => pairs
                .len()
                .checked_mul(2)
                .and_then(|values| values.checked_mul(CheckpointDType::U32.byte_width()))
                .ok_or(CheckpointError::SizeOverflow {
                    context: "BPE pair payload",
                }),
            Self::Float64(values) => values
                .len()
                .checked_mul(CheckpointDType::F64.byte_width())
                .ok_or(CheckpointError::SizeOverflow {
                    context: "f64 tensor payload",
                }),
        }
    }

    fn write_to(self, bytes: &mut Vec<u8>) {
        match self {
            Self::Bytes(values) => bytes.extend_from_slice(values),
            Self::BpePairs(pairs) => {
                for pair in pairs {
                    put_u32(bytes, pair.left());
                    put_u32(bytes, pair.right());
                }
            }
            Self::Float64(values) => {
                for &value in values {
                    put_f64(bytes, value);
                }
            }
        }
    }
}

/// Owns record metadata while borrowing the values that will become file bytes.
#[derive(Debug)]
struct TensorRecordPlan<'a> {
    descriptor: CheckpointTensorDescriptor,
    payload: TensorPayload<'a>,
}

    /// Encodes one canonical little-endian file without native-memory casts.
    pub fn encode(&self) -> Result<EncodedCheckpoint, CheckpointError> {
        let mut records = self.tensor_record_plan()?;
        let model_parameter_count = self.model_state.parameter_names().len();
        let optimizer_state_count = self.optimizer_state.parameter_names().len();

        let mut provisional_header = Vec::new();
        write_fixed_header(&mut provisional_header, 0, 0, 0);
        write_metadata(
            &mut provisional_header,
            self,
            model_parameter_count,
            optimizer_state_count,
            &records,
        )?;
        let header_bytes = usize_to_u64(provisional_header.len(), "header length")?;

        let mut next_offset = header_bytes;
        let mut payload_bytes = 0_u64;
        for record in &mut records {
            record.descriptor.offset = next_offset;
            next_offset = next_offset.checked_add(record.descriptor.byte_len).ok_or(
                CheckpointError::SizeOverflow {
                    context: "tensor end offset",
                },
            )?;
            payload_bytes = payload_bytes
                .checked_add(record.descriptor.byte_len)
                .ok_or(CheckpointError::SizeOverflow {
                    context: "payload length",
                })?;
        }

        let total_bytes =
            header_bytes
                .checked_add(payload_bytes)
                .ok_or(CheckpointError::SizeOverflow {
                    context: "complete file length",
                })?;
        let total_capacity = u64_to_usize(total_bytes, "complete file length")?;
        let mut bytes = Vec::new();
        bytes
            .try_reserve_exact(total_capacity)
            .map_err(|_| CheckpointError::Allocation {
                context: "complete checkpoint file",
            })?;
        write_fixed_header(&mut bytes, header_bytes, payload_bytes, 0);
        write_metadata(
            &mut bytes,
            self,
            model_parameter_count,
            optimizer_state_count,
            &records,
        )?;
        if bytes.len() != u64_to_usize(header_bytes, "header length")? {
            return Err(CheckpointError::Layout(
                "two-pass header length changed".to_owned(),
            ));
        }
        for record in &records {
            record.payload.write_to(&mut bytes);
        }
        if bytes.len() != total_capacity {
            return Err(CheckpointError::Layout(
                "encoded file length changed after layout".to_owned(),
            ));
        }
        let checksum = checkpoint_checksum(&bytes);
        bytes[CHECKSUM_OFFSET..CHECKSUM_OFFSET + CHECKSUM_WIDTH]
            .copy_from_slice(&checksum.to_le_bytes());

        Ok(EncodedCheckpoint {
            bytes,
            header_bytes,
            checksum,
            tensors: records
                .into_iter()
                .map(|record| record.descriptor)
                .collect(),
        })
    }

impl DecoderParameterSource for DecoderModelState {
    fn len(&self) -> usize {
        self.parameters.len()
    }

    fn name(&self, index: usize) -> &str {
        &self.parameters[index].name
    }

    fn with_tensor<R>(&self, index: usize, inspect: impl FnOnce(&Tensor) -> R) -> R {
        inspect(&self.parameters[index].value)
    }
}

impl DecoderModelState {
    /// Builds graph-free state after the caller has validated every parameter leaf.
    pub(crate) fn try_from_leaf_validated_parameters(
        config: DecoderModelConfig,
        parameters: Vec<(String, Tensor)>,
    ) -> Result<Self, TrainerError> {
        let state = Self {
            config,
            parameters: parameters
                .into_iter()
                .map(|(name, value)| StateParameter { name, value })
                .collect(),
        };
        validate_parameter_layout(config, &state)?;
        Ok(state)
    }
}

    /// Rejects the complete file before exposing any partially restored state.
    pub fn from_bytes(bytes: &[u8]) -> Result<Self, CheckpointError> {
        if bytes.len() < FIXED_HEADER_BYTES {
            return Err(CheckpointError::Truncated {
                context: "fixed header",
            });
        }
        if bytes[..CHECKPOINT_MAGIC.len()] != CHECKPOINT_MAGIC {
            return Err(CheckpointError::InvalidMagic);
        }
        let version = u16::from_le_bytes(
            bytes[8..10]
                .try_into()
                .expect("the fixed header length was checked"),
        );
        if version != CHECKPOINT_VERSION {
            return Err(CheckpointError::UnsupportedVersion { found: version });
        }
        let endian = u32::from_le_bytes(
            bytes[10..14]
                .try_into()
                .expect("the fixed header length was checked"),
        );
        if endian != LITTLE_ENDIAN_MARKER {
            return Err(CheckpointError::UnsupportedEndianness { found: endian });
        }
        let header_bytes = read_fixed_u64(bytes, 14);
        let payload_bytes = read_fixed_u64(bytes, 22);
        let stored_checksum = read_fixed_u64(bytes, CHECKSUM_OFFSET);
        let declared_total = header_bytes.checked_add(payload_bytes);
        if header_bytes < usize_to_u64(FIXED_HEADER_BYTES, "fixed header length")?
            || declared_total != Some(usize_to_u64(bytes.len(), "file length")?)
        {
            return Err(CheckpointError::FileExtent {
                header: header_bytes,
                payload: payload_bytes,
                actual: bytes.len(),
            });
        }
        let header_end = u64_to_usize(header_bytes, "header length")?;
        if header_end > bytes.len() {
            return Err(CheckpointError::FileExtent {
                header: header_bytes,
                payload: payload_bytes,
                actual: bytes.len(),
            });
        }
        let actual_checksum = checkpoint_checksum(bytes);
        if stored_checksum != actual_checksum {
            return Err(CheckpointError::ChecksumMismatch {
                expected: stored_checksum,
                actual: actual_checksum,
            });
        }

        let header = decode_variable_header(&bytes[FIXED_HEADER_BYTES..header_end])?;
        validate_descriptors(&header, header_bytes, bytes.len())?;
        let tokenizer_count = match header.tokenizer_kind {
            1 => header.tokenizer_vocabulary,
            2 => 1,
            tag => return Err(CheckpointError::UnknownTokenizerKind { tag }),
        };
        let tokenizer_end = tokenizer_count;
        let model_end = tokenizer_end
            .checked_add(header.model_parameter_count)
            .ok_or(CheckpointError::SizeOverflow {
                context: "model descriptor boundary",
            })?;
        let optimizer_descriptor_count =
            header
                .optimizer_state_count
                .checked_mul(2)
                .ok_or(CheckpointError::SizeOverflow {
                    context: "optimizer descriptor count",
                })?;
        let expected_total = model_end.checked_add(optimizer_descriptor_count).ok_or(
            CheckpointError::SizeOverflow {
                context: "optimizer descriptor boundary",
            },
        )?;
        if header.descriptors.len() != expected_total {
            return Err(CheckpointError::Layout(format!(
                "descriptor table has {} records, expected {expected_total}",
                header.descriptors.len()
            )));
        }

        let tokenizer = decode_tokenizer(
            header.tokenizer_kind,
            header.tokenizer_vocabulary,
            &header.descriptors[..tokenizer_end],
            bytes,
        )?;
        let model_state = decode_model(
            header.config,
            &header.descriptors[tokenizer_end..model_end],
            bytes,
        )?;
        let optimizer_state = decode_optimizer(
            header.optimizer_config,
            header.optimizer_groups,
            header.optimizer_step,
            header.beta1_power,
            header.beta2_power,
            &header.descriptors[model_end..],
            bytes,
        )?;
        let checkpoint = Self::from_owned_parts(
            tokenizer,
            model_state,
            optimizer_state,
            header.selected_step,
            header.sampling_rng_state,
        )?;
        let canonical = checkpoint.encode()?;
        if canonical.bytes() != bytes {
            return Err(CheckpointError::NonCanonical(
                "decoded state does not reproduce the original bytes".to_owned(),
            ));
        }
        Ok(checkpoint)
    }

    /// Replaces one file through a synchronized same-directory temporary name.
    ///
    /// The course's supported Unix workflow relies on same-filesystem rename
    /// semantics. Other targets return an explicit error instead of deleting an
    /// existing destination before a non-atomic move.
    pub fn save_atomic(
        &self,
        path: impl AsRef<Path>,
    ) -> Result<EncodedCheckpoint, CheckpointError> {
        #[cfg(not(unix))]
        {
            let _ = path;
            return Err(CheckpointError::UnsupportedAtomicReplacement);
        }

        #[cfg(unix)]
        {
            let path = path.as_ref();
            let file_name = path.file_name().ok_or_else(|| {
                CheckpointError::Layout("checkpoint destination has no file name".to_owned())
            })?;
            let parent = path
                .parent()
                .filter(|parent| !parent.as_os_str().is_empty())
                .unwrap_or_else(|| Path::new("."));
            let encoded = self.encode()?;
            let (temporary_path, mut temporary) =
                create_temporary(parent, file_name.to_string_lossy().as_ref())?;
            let publication = (|| -> Result<(), CheckpointError> {
                temporary
                    .write_all(encoded.bytes())
                    .map_err(|source| CheckpointError::Io {
                        operation: "write temporary",
                        path: temporary_path.clone(),
                        source,
                    })?;
                temporary.sync_all().map_err(|source| CheckpointError::Io {
                    operation: "synchronize temporary",
                    path: temporary_path.clone(),
                    source,
                })?;
                drop(temporary);
                fs::rename(&temporary_path, path).map_err(|source| CheckpointError::Io {
                    operation: "atomically replace",
                    path: path.to_owned(),
                    source,
                })?;
                File::open(parent)
                    .and_then(|directory| directory.sync_all())
                    .map_err(|source| CheckpointError::Io {
                        operation: "synchronize parent directory",
                        path: parent.to_owned(),
                        source,
                    })?;
                Ok(())
            })();
            if publication.is_err() {
                let _ = fs::remove_file(&temporary_path);
            }
            publication?;
            Ok(encoded)
        }
    }

pub fn learner_evidence() -> Result<LearnerEvidence, FixtureError> {
    let selected = selection_evidence()?;
    let selected_training_state = selected.result.selected_training_state();
    let selected_step = u64::try_from(selected.result.selected_step())
        .map_err(|_| FixtureError::Invariant("selected step does not fit u64"))?;
    require(
        selected_step == selected_training_state.optimizer_state().step_count(),
        "trainer-issued selected model and optimizer no longer share one step",
    )?;

    let mut sampling_rng = SplitMix64::from_seed(SAMPLING_RNG_SEED);
    let _ = sampling_rng.next_u64();
    let saved_sampling_rng_state = sampling_rng.state();
    let expected_sampling_rng_next = sampling_rng.next_u64();
    let checkpoint = Checkpoint::from_snapshot(
        literal_tokenizer()?,
        selected_training_state,
        saved_sampling_rng_state,
    )?;
    let encoded = checkpoint.encode()?;
    let repeated = checkpoint.encode()?;
    let loaded = Checkpoint::from_bytes(encoded.bytes())?;
    let loaded_encoded = loaded.encode()?;

    let original_model = checkpoint.restore_independent_model()?;
    let loaded_optimizer = loaded.restore_optimizer();
    let loaded_sampling_rng_state = loaded.sampling_rng_state();
    let loaded_model = loaded.into_model()?;
    let original_logits = logits_bits(&original_model)?;
    let loaded_logits = logits_bits(&loaded_model)?;
    let (updated_original, updated_original_optimizer) = apply_component_update(
        original_model,
        checkpoint.restore_optimizer(),
        &LOGIT_INPUTS,
        &LOGIT_TARGETS,
        NEXT_LEARNING_RATE,
    )?;
    let (updated_loaded, updated_loaded_optimizer) = apply_component_update(
        loaded_model,
        loaded_optimizer,
        &LOGIT_INPUTS,
        &LOGIT_TARGETS,
        NEXT_LEARNING_RATE,
    )?;
    let updated_original_bits = parameter_bits(&updated_original);
    let updated_loaded_bits = parameter_bits(&updated_loaded);
    let updated_original_logits = logits_bits(&updated_original)?;
    let updated_loaded_logits = logits_bits(&updated_loaded)?;
    let (changed_batch_model, changed_batch_optimizer) = apply_component_update(
        checkpoint.restore_independent_model()?,
        checkpoint.restore_optimizer(),
        &[1, 2],
        &[2, 3],
        NEXT_LEARNING_RATE,
    )?;
    let changed_batch_diverges = parameter_bits(&changed_batch_model) != updated_original_bits
        || changed_batch_optimizer != updated_original_optimizer;
    let (changed_learning_rate_model, changed_learning_rate_optimizer) = apply_component_update(
        checkpoint.restore_independent_model()?,
        checkpoint.restore_optimizer(),
        &LOGIT_INPUTS,
        &LOGIT_TARGETS,
        NEXT_LEARNING_RATE * 2.0,
    )?;
    let changed_learning_rate_diverges = parameter_bits(&changed_learning_rate_model)
        != updated_original_bits
        || changed_learning_rate_optimizer != updated_original_optimizer;
    let loaded_sampling_rng_next = SplitMix64::from_state(loaded_sampling_rng_state).next_u64();
    let corruption = corruption_evidence(selected_training_state, &checkpoint, encoded.bytes())?;
    let atomic = atomic_evidence(selected_training_state, &checkpoint)?;

    require(
        encoded.tensors().len() == 38,
        "checkpoint record count changed",
    )?;
    require(
        checkpoint.model_state().parameter_names().len() == 11,
        "model parameter tensor count changed",
    )?;
    require(
        checkpoint.model_state().scalar_count() == 144,
        "model scalar count changed",
    )?;
    require(
        updated_original_optimizer.step_count() == selected_step + 1,
        "component-replay optimizer step count changed",
    )?;
    require(
        changed_batch_diverges && changed_learning_rate_diverges,
        "caller-supplied batch or learning-rate adversary no longer diverges",
    )?;
    require(
        corruption.version_rejected
            && corruption.vocabulary_mismatch_rejected
            && corruption.step_mismatch_rejected
            && corruption.truncation_rejected
            && corruption.checksum_rejected,
        "one corruption fixture was accepted",
    )?;
    require(
        atomic.replaced_complete_file && atomic.temporary_files == 0,
        "atomic replacement evidence changed",
    )?;
    let history = historical_checkpoint_contrast(&checkpoint, &encoded);
    require(
        history.isolated_parameter_bytes.len()
            == history.isolated_parameter_scalars * std::mem::size_of::<f64>(),
        "isolated parameter byte length changed",
    )?;
    require(
        history.tokenizer_records
            + history.isolated_parameter_tensors
            + history.optimizer_moment_records
            == history.checkpoint_records,
        "checkpoint record-family counts changed",
    )?;

    Ok(LearnerEvidence {
        checkpoint,
        bytes_deterministic: encoded == repeated,
        round_trip_identical: encoded == loaded_encoded,
        logits_before_bits_identical: original_logits == loaded_logits,
        logits_before_fingerprint: bits_fingerprint(&loaded_logits),
        parameter_bits_after_identical: updated_original_bits == updated_loaded_bits,
        optimizer_after_identical: updated_original_optimizer == updated_loaded_optimizer,
        logits_after_bits_identical: updated_original_logits == updated_loaded_logits,
        logits_after_fingerprint: bits_fingerprint(&updated_loaded_logits),
        changed_batch_diverges,
        changed_learning_rate_diverges,
        sampling_rng_next_identical: expected_sampling_rng_next == loaded_sampling_rng_next,
        sampling_rng_next: loaded_sampling_rng_next,
        corruption,
        atomic,
        history,
        encoded,
    })
}

fn main() -> Result<(), Box<dyn std::error::Error>> {
    print!("{}", ch35_checkpoints::learner_report()?);
    Ok(())
}