Traingen parity

def compare_batch_stream(
    reference_loader: Iterable[Mapping[str, Shaped[torch.Tensor, "..."]]],
    target_loader: Iterable[Mapping[str, Shaped[torch.Tensor, "..."]]],
    *,
    steps: int,
) -> BatchStreamReport:
    """Compare two dataloader streams for exact tensor equality."""
    for step, (reference_batch, target_batch) in enumerate(
        zip(reference_loader, target_loader, strict=False)
    ):
        if step >= steps:
            break
        for key, reference_value in reference_batch.items():
            target_value = target_batch.get(key)
            if isinstance(reference_value, torch.Tensor) and isinstance(
                target_value, torch.Tensor
            ):
                if not torch.equal(reference_value, target_value):
                    return BatchStreamReport(False, step + 1, f"{step}:{key}")
    return BatchStreamReport(True, steps)

def compare_optimizer_step(
    reference_state: Mapping[str, Shaped[torch.Tensor, "..."]],
    target_state: Mapping[str, Shaped[torch.Tensor, "..."]],
    tolerances: Mapping[str, TensorTolerance] | None = None,
) -> OptimizerStepReport:
    """Compare two state dictionaries after an optimizer step."""
    tol_map = tolerances or {}
    missing = tuple(name for name in reference_state if name not in target_state)
    comparisons = tuple(
        compare_tensors(
            name, target_state[name], reference_state[name], tol_map.get(name)
        )
        for name in reference_state
        if name not in missing
    )
    return OptimizerStepReport(
        passed=not missing and all(item.passed for item in comparisons),
        comparisons=comparisons,
        missing=missing,
    )

def compare_step_trace(
    reference: StepTrace,
    target: StepTrace,
    tolerances: Mapping[str, TensorTolerance] | None = None,
) -> StepReport:
    """Compare two named step traces.

    Args:
        reference: Reference trace.
        target: Converted implementation trace.
        tolerances: Per-tensor tolerance map.

    Returns:
        Step comparison report.

    Raises:
        RuntimeError: If tensor comparisons fail unexpectedly.

    Examples:
        >>> from traingen_parity.trace import build_step_trace
        >>> a = build_step_trace("a", {"x": torch.ones(1)})
        >>> compare_step_trace(a, a).passed
        True
    """
    tol_map = tolerances or {}
    names = tuple(reference.tensors.keys())
    missing = tuple(name for name in names if name not in target.tensors)
    comparisons = tuple(
        compare_tensors(
            name,
            target.tensors[name],
            reference.tensors[name],
            tol_map.get(name),
        )
        for name in names
        if name not in missing
    )
    return StepReport(
        passed=not missing and all(item.passed for item in comparisons),
        comparisons=comparisons,
        missing=missing,
    )

def compare_tensors(
    name: str,
    actual: Shaped[torch.Tensor, "..."],
    expected: Shaped[torch.Tensor, "..."],
    tolerance: TensorTolerance | None = None,
) -> TensorComparison:
    """Compare two tensors and return max-difference diagnostics.

    Args:
        name: Tensor name.
        actual: Actual tensor.
        expected: Expected tensor.
        tolerance: Absolute and relative tolerance. Defaults to exact equality.

    Returns:
        Tensor comparison report.

    Raises:
        RuntimeError: If tensors cannot be broadcast for comparison.

    Examples:
        >>> compare_tensors("x", torch.ones(1), torch.ones(1)).passed
        True
    """
    tol = tolerance or TensorTolerance()
    if actual.shape != expected.shape:
        return TensorComparison(
            name, False, float("inf"), float("inf"), "shape mismatch"
        )
    actual_detached = actual.detach()
    expected_detached = expected.detach()
    diff = (actual_detached.float() - expected_detached.float()).abs()
    max_abs = float(diff.max().item()) if diff.numel() else 0.0
    denom = expected_detached.float().abs().clamp_min(torch.finfo(diff.dtype).eps)
    rel = diff / denom
    max_rel = float(rel.max().item()) if rel.numel() else 0.0
    if actual_detached.is_floating_point() or expected_detached.is_floating_point():
        passed = torch.allclose(actual, expected, atol=tol.atol, rtol=tol.rtol)
    else:
        passed = torch.equal(actual, expected)
    message = "ok" if passed else f"max_abs={max_abs:.6g}, max_rel={max_rel:.6g}"
    return TensorComparison(name, passed, max_abs, max_rel, message)

def apply_determinism(config: DeterminismConfig) -> None:
    """Apply deterministic runtime settings.

    Args:
        config: Determinism configuration.

    Returns:
        None.

    Raises:
        RuntimeError: If torch cannot enable deterministic algorithms.

    Examples:
        >>> apply_determinism(DeterminismConfig(deterministic_algorithms=False))
    """
    if config.cublas_workspace_config is not None:
        os.environ.setdefault("CUBLAS_WORKSPACE_CONFIG", config.cublas_workspace_config)
    random.seed(config.seed)
    np.random.seed(config.seed)
    torch.manual_seed(config.seed)
    if torch.cuda.is_available():
        torch.cuda.manual_seed_all(config.seed)
    torch.backends.cuda.matmul.allow_tf32 = config.allow_tf32
    torch.backends.cudnn.allow_tf32 = config.allow_tf32
    torch.backends.cudnn.benchmark = config.cudnn_benchmark
    torch.use_deterministic_algorithms(config.deterministic_algorithms)

def capture_rng_state() -> RNGState:
    """Capture all RNG states needed for step-level parity.

    Args:
        None.

    Returns:
        RNG state dataclass.

    Raises:
        RuntimeError: If torch cannot read CUDA RNG state.

    Examples:
        >>> state = capture_rng_state()
        >>> isinstance(state.torch_cpu, torch.Tensor)
        True
    """
    cuda_state = (
        tuple(torch.cuda.get_rng_state_all()) if torch.cuda.is_available() else ()
    )
    return RNGState(
        python=cast(PythonRandomState, random.getstate()),
        numpy=cast(
            tuple[str, UInt32[np.ndarray, "state"], int, int, float],
            np.random.get_state(),
        ),
        torch_cpu=torch.random.get_rng_state(),
        torch_cuda=cuda_state,
    )

def restore_rng_state(state: RNGState) -> None:
    """Restore a state captured with :func:`capture_rng_state`.

    Args:
        state: Previously captured RNG state.

    Returns:
        None.

    Raises:
        RuntimeError: If CUDA state restoration fails.

    Examples:
        >>> state = capture_rng_state()
        >>> restore_rng_state(state)
    """
    random.setstate(state.python)
    np.random.set_state(state.numpy)
    torch.random.set_rng_state(state.torch_cpu)
    if torch.cuda.is_available() and state.torch_cuda:
        torch.cuda.set_rng_state_all(list(state.torch_cuda))

def build_step_trace(
    name: str,
    tensors: dict[str, Shaped[torch.Tensor, "..."]],
    *,
    metadata: TraceMetadata | None = None,
) -> StepTrace:
    """Build a trace from named tensors.

    Args:
        name: Trace name.
        tensors: Named tensor values.
        metadata: Optional non-tensor metadata.

    Returns:
        Step trace with summaries.

    Raises:
        RuntimeError: If tensor summaries cannot be computed.

    Examples:
        >>> trace = build_step_trace("step", {"loss": torch.tensor(1.0)})
        >>> "loss" in trace.summaries
        True
    """
    return StepTrace(
        name=name,
        tensors=tensors,
        summaries={key: summarize_tensor(value) for key, value in tensors.items()},
        metadata=metadata or {},
    )

def scalar_trace_value(value: Float[torch.Tensor, ""]) -> float:
    """Return a Python scalar from a scalar tensor."""
    return float(value.detach().cpu().item())

def summarize_tensor(tensor: Shaped[torch.Tensor, "..."]) -> TensorSummary:
    """Build a deterministic tensor summary.

    Args:
        tensor: Tensor to summarize.

    Returns:
        Tensor summary dataclass.

    Raises:
        RuntimeError: If tensor statistics cannot be computed.

    Examples:
        >>> summarize_tensor(torch.ones(2)).mean
        1.0
    """
    detached = tensor.detach()
    stats = detached.float()
    if detached.numel() == 0:
        min_value = max_value = mean_value = None
    else:
        min_value = float(stats.min().item())
        max_value = float(stats.max().item())
        mean_value = float(stats.mean().item())
    return TensorSummary(
        shape=tuple(detached.shape),
        dtype=str(detached.dtype),
        device=str(detached.device),
        sha256=tensor_sha256(detached),
        min=min_value,
        max=max_value,
        mean=mean_value,
    )

def tensor_sha256(tensor: Shaped[torch.Tensor, "..."]) -> str:
    """Return a SHA-256 digest for tensor bytes on CPU.

    Args:
        tensor: Tensor to hash.

    Returns:
        Hex digest of the contiguous CPU tensor bytes.

    Raises:
        RuntimeError: If the tensor cannot be copied to CPU.

    Examples:
        >>> tensor_sha256(torch.tensor([1, 2])).startswith("0")
        False
    """
    array = tensor.detach().contiguous().cpu().numpy()
    return hashlib.sha256(array.tobytes()).hexdigest()

def trace_training_step(
    module: TrainingStepModule,
    batch: dict[str, Shaped[torch.Tensor, "..."]],
    rng_state: RNGState | None,
    trace_points: tuple[str, ...],
) -> StepTrace:
    """Run ``module.training_step`` and collect requested trace points.

    Args:
        module: Object exposing ``training_step`` and optionally
            ``latest_step_trace``.
        batch: Training batch.
        rng_state: Optional RNG state restored before the step.
        trace_points: Requested tensor names.

    Returns:
        Step trace for requested tensor names.

    Raises:
        AttributeError: If the module does not expose ``training_step``.

    Examples:
        >>> class M:
        ...     def training_step(self, batch, batch_idx):
        ...         self.latest_step_trace = {"loss": torch.tensor(1.0)}
        ...         return torch.tensor(1.0)
        >>> trace_training_step(M(), {}, None, ("loss",)).tensors["loss"].item()
        1.0
    """
    if rng_state is not None:
        restore_rng_state(rng_state)
    loss = module.training_step(batch, 0)
    raw = dict(getattr(module, "latest_step_trace", {}))
    raw.setdefault("train_loss", loss)
    tensors = {
        key: value
        for key, value in raw.items()
        if key in trace_points and isinstance(value, torch.Tensor)
    }
    return build_step_trace(
        getattr(module, "__class__", type(module)).__name__,
        tensors,
        metadata={"trace_points": trace_points},
    )