Diff Coverage

        dict: What the file holds, keyed by physical name.
    """
    if ckpt_format == "safetensors":
        return load_file(filename=file_path)
    return torch.load(f=file_path, weights_only=True)

        verbose: Whether the caller asked for debug output.
        my_rank: Current process rank, prefixed to every line.
        message: Stage-specific text; the caller formats it.
    """
    if verbose:
        print(f"Rank {my_rank}: {message}")


def generate_groups_from_template(
    template: Union[list[int], tuple[int, ...]],

    world_size = int(world_size)
    my_rank = int(my_rank)
    template_len = len(template)

    _report_template_stage(verbose, my_rank, f"Original Template = {template}, World size = {world_size}")

    if template_len == 1:
        return [[i] for i in range(world_size)]

    if template_len < 2:
        raise ValueError(f"Template must have at least 2 ranks, got {template}")

    # 1. Template normalization: convert to 0-based template
    normalized_template = [x - template[0] for x in template]
    _report_template_stage(verbose, my_rank, f"Normalized Template = {normalized_template}")

    # 2. Analyze normalized template core params
    # intra-step: spacing between elements in template
    intra_step = _validate_intra_step(normalized_template, template_len)

    # 2. Analyze normalized template core params
    # intra-step: spacing between elements in template
    intra_step = _validate_intra_step(normalized_template, template_len)
    # block size: ranks per block; it is also the spacing between adjacent blocks
    block_size = intra_step * template_len
    template_span = normalized_template[-1] - normalized_template[0]
    _report_template_stage(
        verbose, my_rank,
        f"Template analysis - intra_step={intra_step}, template_span={template_span}, block_size={block_size}"
    )

        f"Template analysis - intra_step={intra_step}, template_span={template_span}, block_size={block_size}"
    )

    # 3. Compute all valid block start positions
    group_starts = _compute_group_starts(world_size, block_size, block_size)
    _report_template_stage(verbose, my_rank, f"Possible block starts: {group_starts}")

    # 4. Generate all valid sub-groups for each block
    all_groups = _build_groups_for_blocks(
        group_starts, block_size, template_span,

    # 6. Sort: ensure all processes generate groups in same order
    all_groups.sort(key=lambda x: (x[0], x[1] if len(x) > 1 else 0))

    _report_template_stage(
        verbose, my_rank,
        f"Generated {len(all_groups)} groups, covering {len(unique_ranks)} unique ranks\n"
        f"Final group list: {all_groups}"
    )

        Safe to call at any point, including while the context is active.
        The observer/hook installation is left untouched — only the data.
        """
        self._reset_records()

    def _reset_records(self):
        """Drop records, counters and module info."""
        self._call_stack.clear()

    def _on_op_dispatch_enter(self, op_name: str, op_call, args, kwargs):  # pylint: disable=W0613
        """Called by OpDispatcher.dispatch() before the op executes."""
        if not self._active:
            return

        record = self._make_record(
            OpCall,
            op_name=op_name,

            op_name=op_name,
            input_infos=self._extract_tensor_infos(args),
        )
        if record is None:
            return

        self._attach(record)

    def _on_op_dispatch_exit(self, op_name, result):  # pylint: disable=W0613

        # Peek before popping: a nested scope that never saw an enter must not
        # pop someone else's frame off the stack.
        if not isinstance(self._call_stack[-1], OpCall):
            return

        record = self._call_stack.pop()
        record.output_infos = self._extract_tensor_infos((result,))

            in that case.
        """
        depth = len(self._call_stack)
        if depth >= _MAX_DEPTH or self._record_count >= _MAX_RECORDS:
            self._dropped_records += 1
            return None

        self._record_count += 1
        return record_cls(call_depth=depth, **fields)

        self._call_stack.append(record)

    def get_dropped_record_count(self) -> int:
        """Returns how many records were discarded because the budget ran out."""
        return self._dropped_records

    # ------------------------------------------------------------------
    # Collective tracer callback
    # ------------------------------------------------------------------

        # being dropped, so it is updated before anything else.
        self._comm_counts[method_name] += 1

        if not self._active:
            return

        input_shape = None
        input_dtype = ""
        if args and hasattr(args[0], "shape"):

            output_shape=output_shape,
            input_dtype=input_dtype,
        )
        if record is None:
            return

        # A collective is a leaf: it is linked but never pushed, so it cannot
        # swallow the enclosing op's exit.
        if self._call_stack:

    # ------------------------------------------------------------------

    def _on_module_event(self, module_fqn: str, event_type: str):
        """Invoked by ModuleTracker on forward enter/exit."""
        if not self._active:
            return

        if event_type == "exit":
            # Only unwind a frame this tracer opened.
            if self._call_stack and isinstance(self._call_stack[-1], AnnotateCall):
                self._call_stack.pop()
            return

        record = self._make_record(
            AnnotateCall,
            module_fqn=module_fqn,
            event_type=event_type,
        )
        if record is None:
            return

        self._attach(record)

    # ------------------------------------------------------------------
    # Module info collection
    # ------------------------------------------------------------------

        if noise_level is None:
            noise_level = 1

        if self._dropped_records:
            logger.warning(
                "%d record(s) were dropped because the record budget was exhausted; "
                "the table is incomplete. Only trace a bounded number of steps.",
                self._dropped_records,
            )

            "total_counts": self.get_total_counts(),
            "records": [],
        }

        if self._dropped_records:
            data["dropped_records"] = self._dropped_records

        if self._sharding_info:
            data["sharding_info"] = {k: str(v) for k, v in self._sharding_info.items()}

            lines.append(f"{prefix}{'Collective':<20} {record.render_self()}")
        elif isinstance(record, OpCall) and noise_level >= 1:
            lines.append(f"{prefix}{'Op':<20} {record.render_self()}")
        elif isinstance(record, AnnotateCall) and noise_level >= 2:
            lines.append(f"{prefix}{'Module':<20} {record.render_self()}")

        for child in record.children:
            self._collect_table_lines(child, lines, noise_level, indent + 1)