Diff Coverage

def tile_box(offsets: tuple[int, ...], lengths: tuple[int, ...], limit: int) -> Iterator[Box]:
    """Split a rectangle into non-overlapping tiles of at most ``limit`` elements."""
    if limit < 1:
        raise ValueError("Tile limit must be positive")
    if any(size == 0 for size in lengths):
        return
    if math.prod(lengths) <= limit:
        yield offsets, lengths

    low = [0] * len(block.lengths)
    high = list(block.lengths)
    for size, column in zip(block.lengths, block.steps):
        if size > 1 and sum(step != 0 for step in column) != 1:
            raise ValueError("Inverse conversion has a non-rectangular shard mapping")
    for dim, origin in enumerate(block.base):
        steps = tuple(column[dim] for column in block.steps)
        moving = [(axis, step) for axis, step in enumerate(steps) if step]
        start, stop = shard_offsets[dim], shard_offsets[dim] + shard_shape[dim]

            raise RuntimeError("HFLoadPlanner.finish() needs a load planned with this planner first")
        for target in self._targets.values():
            padded = self._padded_targets.get(id(target))
            if padded is not None:
                self._zero_padding(padded[0], target, padded[1])
            if id(target) in self._loaded_ids:
                target._is_hf_initialized = True  # pylint: disable=protected-access
        self.apply_used_conversions(self.model)
        validate_load_result(self.report.missing_keys, self.report.unexpected_keys, strict=False)

    @staticmethod
    def _zero_padding(name: str, target: torch.Tensor, logical_shape: tuple[int, ...]) -> None:
        """Zero local padding beyond the declared logical checkpoint shape."""
        view = _dcp_view(target)
        chunks = create_chunk_list_for_tensor(view)
        if not chunks:
            return
        if len(chunks) != 1:
            raise ValueError(f"Expected one rectangular local padded parameter shard for {name}")
        chunk = chunks[0]
        local = local_target_tensor(target)
        if tuple(local.shape) != tuple(chunk.sizes):
            raise ValueError(f"Local padded parameter storage does not match its shard metadata: {name}")
        with torch.no_grad():
            for offsets, lengths in checkpoint_padding_boxes(logical_shape, tuple(chunk.offsets), tuple(local.shape)):
                slices = tuple(slice(start, start + size) for start, size in zip(offsets, lengths))
                local[slices].zero_()

    def apply_used_conversions(self, model: nn.Module) -> None:
        """
        Record on ``model`` the conversions the load used, where saving in the original format finds them.

            key = self._state_keys[target_id]
            logical_shape = self._logical_shapes.get(name)
            self._check_target(name, target, self._load_state[key], value, logical_shape)
            if logical_shape is not None and tuple(value.shape) == logical_shape:
                if tuple(self._load_state[key].shape) != logical_shape:
                    self._padded_targets[target_id] = (name, logical_shape)
            if isinstance(value, RegionTensor):
                plan.table[key] = value.remap_blocks()
            else:
                plan.deferred.setdefault(id(value.conversion), (value.conversion, []))[1].append((value.name, name))

    for module_name, module in model.named_modules():
        provider = getattr(module, "checkpoint_logical_shapes", None)
        if not callable(provider):
            continue
        declared = provider()
        if not isinstance(declared, dict):
            raise ValueError(f"Invalid checkpoint padding declaration on {module_name}")
        for parameter_name, shape in declared.items():
            parameter = module.get_parameter(parameter_name)
            if (not isinstance(shape, tuple) or len(shape) != parameter.ndim
                    or any(not isinstance(size, int) or isinstance(size, bool) or size < 0 for size in shape)):
                raise ValueError(f"Invalid logical checkpoint shape on {module_name}: {shape}")
            name = join_fqn(module_name, parameter_name)
            if name in shapes and shapes[name] != shape:
                raise ValueError(f"Conflicting logical checkpoint shapes for {name}")
            shapes[name] = shape
    return shapes


def checkpoint_padding_boxes(

    Returns:
        Local offsets and dimensions of non-overlapping padding boxes.
    """
    if len(logical_shape) != len(offsets) or len(offsets) != len(sizes):
        raise ValueError("Logical shape and local shard must have the same number of dimensions")
    valid = tuple(max(0, min(size, logical - offset))
                  for logical, offset, size in zip(logical_shape, offsets, sizes))
    boxes = []
    for axis, (length, size) in enumerate(zip(valid, sizes)):

    """Copy a checkpoint tensor into its materialized local target."""
    layout = target_layout(target)
    target_shape = tuple(getattr(layout, "tensor_shape", None) or target.shape)
    if logical_shape is not None and tuple(full_tensor.shape) == logical_shape and target_shape != logical_shape:
        if (len(logical_shape) != len(target_shape)
                or any(logical > physical for logical, physical in zip(logical_shape, target_shape))):
            raise ValueError(f"Invalid checkpoint padding shape for {target_name}")
        padded = full_tensor.new_zeros(target_shape)
        padded[tuple(slice(0, size) for size in logical_shape)].copy_(full_tensor)
        full_tensor = padded
    local_tensor = shard_for_target(target_name, full_tensor, target)
    destination = local_target_tensor(target)
    if destination.is_meta:
        raise ValueError(f"Target must be materialized before loading: {target_name}")

    def __init__(self, path: str) -> None:
        """Read indexed tensor headers and optional quantization metadata."""
        location = Path(path).expanduser().resolve()
        if not location.exists():
            raise ValueError(f"Weight checking requires a local HF checkpoint: {location}")
        self.index = resolve_checkpoint_index(str(location))
        self.quantized = build_quantized_hf_reader(location)
        self.shapes = {}
        if self.quantized is not None:

                with safe_open(str(filename), framework="pt", device="cpu") as reader:
                    for name in names:
                        self.shapes[name] = tuple(reader.get_slice(name).get_shape())
        if not self.shapes:
            raise ValueError("No logical HF tensors were selected for checking")

    def read(self, name: str, offsets: tuple[int, ...], lengths: tuple[int, ...]) -> torch.Tensor:
        """Decode a bounded source rectangle, keeping ordinary source dtypes."""
        if self.quantized is not None and len(lengths) == 2:

        if self.quantized is not None and len(lengths) == 2:
            return self.quantized.read_region(name, offsets, lengths)
        with safe_open(str(self.index.files_by_key[name]), framework="pt", device="cpu") as reader:
            if not lengths:
                return reader.get_tensor(name)
            region = tuple(slice(start, start + count) for start, count in zip(offsets, lengths))
            return reader.get_slice(name)[region]

        self.config = model.config
        self._weight_conversions = conversions

    def __getattr__(self, name: str) -> Any:
        return getattr(self.model, name)


def _model_shards(model: nn.Module) -> dict[str, CheckpointTensorShard]:
    """Describe every materialized local model shard for comparison."""

        if not callable(provider):
            raise ValueError("External state must expose checkpoint_tensor_shards() for complete weight checking")
        external_shards = provider()
        if set(external_shards) != set(external.parameters):
            raise ValueError("External checkpoint shard descriptors do not cover exactly the external parameters")
    shards = {}
    for name, target in build_load_targets(model).items():
        if name in external_shards:
            shards[name] = external_shards[name]

            continue
        layout = target_layout(target)
        view = target
        if layout is not None:
            reject_partial_layout(name, layout)
            if not isinstance(view, DTensor):
                shape = getattr(layout, "tensor_shape", None)
                view = DTensor.from_local_with_layout(target.detach(), layout, shape=shape)
        chunks = create_chunk_list_for_tensor(view)
        if len(chunks) != 1:
            raise ValueError(f"Weight checking requires one rectangular local shard: {name}")
        local = local_target_tensor(target)
        if local.is_meta or tuple(local.shape) != tuple(chunks[0].sizes):
            raise ValueError(f"Live storage does not match its checkpoint shard metadata: {name}")
        logical_shape = logical_shapes.get(name)
        if logical_shape == tuple(view.shape):
            logical_shape = None
        shards[name] = CheckpointTensorShard(
            local, tuple(view.shape), tuple(chunks[0].offsets), logical_shape,
        )
    if not shards:
        raise ValueError("The model has no checkpoint tensors to check")
    return shards


def _inverse_regions(model: nn.Module, shards: dict[str, CheckpointTensorShard]) -> dict[str, RegionTensor]:

        stages = [replacements, getattr(model, "_hp_used_base_weight_conversions", [])]
    else:
        conversions = getattr(model, "_weight_conversions", None)
        if conversions is None:
            conversions = get_model_conversion_mapping(model, add_legacy=False)
        stages = [conversions]
    for conversions in stages:
        converters = [item for item in conversions if isinstance(item, WeightConverter)]
        renamings = [item for item in conversions if isinstance(item, WeightRenaming)]

    for conversions in stages:
        converters = [item for item in conversions if isinstance(item, WeightConverter)]
        renamings = [item for item in conversions if isinstance(item, WeightRenaming)]
        if len(converters) + len(renamings) != len(conversions):
            raise ValueError("Unsupported recorded weight-conversion type")
        if converters:
            regions = _reverse_stage(model, regions, converters)
        for renaming in reversed(renamings):
            regions = _reverse_stage(model, regions, [renaming])

            regions = _reverse_stage(model, regions, converters)
        for renaming in reversed(renamings):
            regions = _reverse_stage(model, regions, [renaming])
    if any(not isinstance(value, RegionTensor) for value in regions.values()):
        raise ValueError("Inverse conversion produced non-symbolic weights; bounded verification is unsupported")
    return regions


def _reverse_stage(model: nn.Module, regions: dict[str, RegionTensor],

            relative = name[len(prefix):]
            if any(probe.rename_source_key(relative)[1] is not None for probe in probes):
                lookup = relative
        if lookup in normalized:
            raise ValueError(f"Ambiguous inverse conversion after resolving base-model prefix: {lookup}")
        normalized[lookup] = region
    return revert_weight_conversion(_InverseModel(model, conversions), normalized)

        (origin - start) * stride for origin, start, stride in zip(base, shard.offsets, strides)
    )
    values = tensor.as_strided(lengths, tile_strides, storage_offset).cpu()
    for operation in block.post:
        values = operation(values)
    return values


def _difference(actual: torch.Tensor, expected: torch.Tensor, storage_dtype: torch.dtype) -> dict[str, Any]:

def _difference(actual: torch.Tensor, expected: torch.Tensor, storage_dtype: torch.dtype) -> dict[str, Any]:
    """Count raw and load-dtype mismatches for one pair of tensor tiles."""
    if actual.is_complex() or expected.is_complex():
        raise ValueError("Complex weight checking is not supported")
    if tuple(actual.shape) != tuple(expected.shape):
        raise ValueError(f"Compared weight tiles have different shapes: {actual.shape}, {expected.shape}")
    raw_bad = actual != expected
    if actual.is_floating_point() and expected.is_floating_point():
        raw_bad |= (actual == 0) & (expected == 0) & (torch.signbit(actual) != torch.signbit(expected))
    nonfinite = ~(torch.isfinite(actual) & torch.isfinite(expected))

    if logical is None or logical == shard.global_shape:
        return None
    if (len(logical) != len(shard.global_shape)
            or any(left > right for left, right in zip(logical, shard.global_shape))):
        raise ValueError(f"Invalid logical checkpoint shape for {name}")
    record = {"model_key": name, "kind": "padding", "elements": 0}
    for box in checkpoint_padding_boxes(logical, shard.offsets, tuple(shard.tensor.shape)):
        for offsets, sizes in tile_box(*box, limit):
            slices = tuple(slice(left, left + size) for left, size in zip(offsets, sizes))

    prefix = getattr(model, "base_model_prefix", "") + "."
    for number, (recovered_name, region) in enumerate(sorted(recovered.items()), start=1):
        name = recovered_name
        if name not in reference.shapes and name.startswith(prefix):
            name = name[len(prefix):]
        if name not in reference.shapes:
            unmatched.append(recovered_name)
            continue
        if tuple(region.shape) != reference.shapes[name]:
            raise ValueError(f"Recovered HF shape mismatch: {name}: {tuple(region.shape)} vs {reference.shapes[name]}")
        record = {
            "hf_key": name, "shape": list(region.shape), "elements": 0,
            "model_keys": sorted(region.sources()),
        }

            "hf_key": name, "shape": list(region.shape), "elements": 0,
            "model_keys": sorted(region.sources()),
        }
        if math.prod(region.shape) >= 1 << 30 or number % 200 == 0:
            logger.info("Checking HF tensor %d/%d: %s, shape=%s", number, len(recovered), name, tuple(region.shape))
        for block in region.blocks:
            shard = shards[block.source]
            box = local_block_box(block, shard.offsets, tuple(shard.tensor.shape))
            if box is None:

def _gather(value: Any) -> list[Any]:
    if not (dist.is_available() and dist.is_initialized()):
        return [value]
    values = [None] * dist.get_world_size()
    dist.all_gather_object(values, value)
    return values


def _summarize(reports: list[dict[str, Any]], shapes: dict[str, tuple[int, ...]]) -> dict[str, Any]:
    """Combine rank reports into global mismatch and coverage results."""

    raw_mismatches = 0
    nonfinite = 0
    for rank, report in enumerate(reports):
        if report.get("error"):
            errors.append({"rank": rank, "error": report["error"]})
        for name, boxes in report.get("source_coverage", {}).items():
            coverage[name].extend(boxes)
        if report.get("uncovered_model_shards") or report.get("unmatched_recovered_keys"):
            errors.append({"rank": rank, "error": "Model weights are not completely covered; see rank report"})

            lines.append(f"  {row['name']}: rank={first['rank']} index={index} "
                         f"actual={first['actual']} reference={first['reference']} "
                         f"reference_in_load_dtype={first['reference_in_load_dtype']}")
    if len(rows) > 50:
        lines.append(f"Showing 50/{len(rows)} rows (mismatches first); complete table: {root / 'tensors.csv'}")
    lines.append("Checked elements and mismatch counts in table include replicated rank observations.")
    lines.append(f"Reports: {root}")
    logger.info("\n%s", "\n".join(lines))

    shapes = {}
    local = {"rank": rank, "phase": phase, "checkpoint": str(pretrained_path)}
    try:
        if chunk_elements < 1:
            raise ValueError("chunk_elements must be positive")
        root.mkdir(parents=True, exist_ok=True)
        reference = _Reference(pretrained_path)
        shapes = reference.shapes
        if rank == 0:

        local.update(_check_local(model, reference, chunk_elements))
        for field in ("exact_reference_mismatches", "load_dtype_mismatches", "nonfinite_elements"):
            local[field] = sum(record.get(field, 0) for record in local["tensors"])
        (root / f"rank_{rank:05d}.json").write_text(json.dumps(local, indent=2, allow_nan=False), encoding="utf-8")
    except Exception as exc:  # pylint: disable=broad-exception-caught
        # A local failure must reach the collective so other ranks cannot continue training.
        local["error"] = f"{type(exc).__name__}: {exc}"
        try:
            root.mkdir(parents=True, exist_ok=True)
            (root / f"rank_{rank:05d}.json").write_text(json.dumps(local, indent=2), encoding="utf-8")
        except OSError:
            logger.exception("Cannot write HF weight-check rank report")
    return local, shapes


def _add_result_counts(rows: list[dict[str, Any]], summary: dict[str, Any]) -> None:

def _write_summary_reports(rank: int, root: Path, rows: list[dict[str, Any]],
                           summary: dict[str, Any]) -> str | None:
    """Write rank-zero aggregate reports and return any filesystem error."""
    if rank != 0:
        return None
    try:
        root.mkdir(parents=True, exist_ok=True)
        _write_table(rows, root)
        (root / "summary.json").write_text(json.dumps(summary, indent=2), encoding="utf-8")

        root.mkdir(parents=True, exist_ok=True)
        _write_table(rows, root)
        (root / "summary.json").write_text(json.dumps(summary, indent=2), encoding="utf-8")
        _print_table(rows, summary, root)
    except OSError as exc:
        return str(exc)
    return None


def _raise_check_failure(summary: dict[str, Any], write_errors: list[Any], root: Path) -> None:

        self.train(module.training)

    def checkpoint_logical_shapes(self) -> dict[str, tuple[int, ...]]:
        """Return the logical embedding shape for checkpoint loading and checking."""
        return {"embed.weight": (self.logical_num_embeddings, int(self.embed.weight.shape[1]))}

    def _fuse(self, hidden_states: torch.Tensor, embeddings: torch.Tensor) -> torch.Tensor:
        """Apply the exact V4.1 WKV projection and signed-root gate."""
        key_value = self.wkv(embeddings.flatten(start_dim=-2))