Diff Coverage

    seq_len = data_transform_raw.get("max_seq_len") or _get_dict(raw, "data").get("max_seq_len")
    if seq_len:
        model_spec["max_position_embeddings"] = seq_len
    else:
        logger.warning(
            "no dataset.data_transform.max_seq_len (nor data.max_seq_len): costing "
            "the model's context limit of %s tokens, which for a long-context model "
            "puts every candidate out of memory",
            model_spec.get("max_position_embeddings", 4096),

            "the model's context limit of %s tokens, which for a long-context model "
            "puts every candidate out of memory",
            model_spec.get("max_position_embeddings", 4096),
        )
        model_spec.setdefault("max_position_embeddings", 4096)
    if context_raw.get("device_num") is not None:
        model_spec["device_num"] = int(context_raw["device_num"])
    model_spec["local_batch_size"] = training_raw.get("micro_batch_size", 1)
    model_spec["compute_dtype"] = model_raw.get("torch_dtype", "bfloat16")

    Returns:
        The pinned degree, else the only candidate, else 1.
    """
    pinned = config.constraint.get(constraint_key)
    if not pinned:
        candidates = config.search_space.get(space_key) or []
        pinned = candidates[0] if len(candidates) == 1 else 1
    return int(pinned)


def _format_result(best_entry: tuple, config: NormalizedConfig) -> Dict[str, Any]:
    """Convert the best ND result entry into a flat result dict."""

        "micro_batch_num": ("fixed_micro_batch_num", "micro_batch_num"),
    }
    for key, (constraint_key, space_key) in fixed_from.items():
        if key not in result:
            result[key] = _pinned_degree(config, constraint_key, space_key)
    total_dp = result.get("dp", 1)
    if "dp_shard" not in result:
        # OP absent from the searched dimensions: fall back to the declared
        # degree, which is what the parser used for the whole run.

    check below with two numbers and no hint of where either came from.
    """
    searched = int(config.constraint.get("global_batch_size", 0) or 0)
    if not searched:
        return
    section = data.get("training") if is_auto_models_schema(data) else data.get("train")
    if not isinstance(section, dict):
        return
    # Only when the strategy really was scored at that batch size; otherwise
    # the check below is the right place for the disagreement to surface.
    resolved = config.resolved_strategy or {}
    implied = (

        * int(resolved.get("dp", 0) or 0)
    )
    if implied != searched:
        return
    current = section.get("global_batch_size")
    if current is not None and int(current) != searched:
        logger.info(
            "global_batch_size %s in the train yaml, %s in the search: "
            "writing the searched value", current, searched,
        )
    section["global_batch_size"] = searched


def write_resolved_yaml(
    config: NormalizedConfig,

    def set(self, mem: "Memory") -> "Memory":
        """Take the size of another memory, in place"""
        if self.__class__ is not mem.__class__:
            return NotImplemented
        self.size, self.unit = mem.size, mem.unit
        return self

    def decrease(self, mem):

            if "linear" not in kind:
                _restore_full_attention(lccfg)
                return
            if not (linear["n_v"] and linear["d_v"] and linear["d_k"]):
                logger.warning(
                    "layer_types declares %s but the config carries no linear "
                    "attention dimensions; costing it as full attention", kind,
                )
                _restore_full_attention(lccfg)
                return
            if getattr(lccfg, "full_attn", None) is None:
                lccfg.full_attn = {
                    name: _declared(lccfg, name) for name in _LINEAR_ATTN_FIELDS
                }

def output_dir() -> str:
    """Return the directory the debug artifacts are written to."""
    if _OUTPUT_DIR:
        return _OUTPUT_DIR
    return os.path.join(os.path.dirname(os.path.abspath(__file__)), "output")


class PerfParts(Enum):

            self.global_batch_size,
            str(self.config.dimensions),
        )
        if self.enable_debug:
            output_path = Debug.output_dir()
            if scored_space:
                Debug.plot_nd(
                    scored_space,
                    output_path,

        search_cfg.constraint["global_batch_size"] = cli_args.global_batch_size
    if cli_args.max_mem is not None:
        # -M sets the device budget the search checks against, exactly as it
        # does on the CLI path, instead of being silently ignored here.
        search_cfg.cluster_spec["device_memory_gb"] = (
            Memory.from_string(cli_args.max_mem.strip()).to_gb().size
        )
    if cli_args.devices is not None:
        cards_per_node = search_cfg.cluster_spec.get("cards_per_node")

        cards_per_node = search_cfg.cluster_spec.get("cards_per_node")
        if not cards_per_node:
            # The device type knows its node size (A3: 16); defaulting to 8
            # silently halves an A3 node and invalidates every candidate.
            device = Hard.device_map.get(cli_args.device_type)
            cards_per_node = device.intra_node_num() if device else 8
            search_cfg.cluster_spec["cards_per_node"] = cards_per_node
            logger.info(
                "cluster.cards_per_node not set, using %d from device type %s",
                cards_per_node, cli_args.device_type,
            )
        cards_per_node = max(1, cards_per_node)

                cards_per_node, cli_args.device_type,
            )
        cards_per_node = max(1, cards_per_node)
        if cli_args.devices % cards_per_node:
            logger.warning(
                "devices=%d is not a multiple of cards_per_node=%d: "
                "%d device(s) will not be placed",
                cli_args.devices, cards_per_node,
                cli_args.devices % cards_per_node,

    search_cfg = read_search_config(cli_args.search_config)
    _apply_cli_overrides(search_cfg, cli_args)

    if getattr(search_cfg, "parallelism_summary", ""):
        logger.output("Parallelism: %s", search_cfg.parallelism_summary)

    errors = validate(search_cfg)
    hard_errors = [e for e in errors if e.severity == "error"]
    warnings = [e for e in errors if e.severity == "warning"]

        layer,
        with_recomp,
    )
    if ccfg.ttype == PerformanceType.TIME:
        flop = estimate_comp_flop_time(lcfg, flop)
    return flop


# Evaluation functions