modules.context

# function-first packer (lift from run_repo_adapter.py)

from typing import List, Optional, Tuple
import os
import torch  # type: ignore

def pack_context_heads(program_root: str, files: List[str], tok, budget_tokens: int) -> str:
    lines_out: List[str] = ["Program snippets:"]
    used = 0
    for rel in files:
        abs_fp = rel if os.path.isabs(rel) else os.path.abspath(os.path.join(program_root, rel))
        try:
            src_lines = open(abs_fp, "r", encoding="utf-8", errors="ignore").read().splitlines()
        except Exception:
            continue
        head_n = min(len(src_lines), 120)
        block = [f"[ctx] path: {os.path.relpath(abs_fp, program_root)}:1-{head_n}"] + src_lines[:head_n] + [""]
        text = "\n".join(block) + "\n"
        t = len(tok(text).input_ids)
        if used + t > budget_tokens:
            head_n = min(len(src_lines), 60)
            block = [f"[ctx] path: {os.path.relpath(abs_fp, program_root)}:1-{head_n}"] + src_lines[:head_n] + [""]
            text = "\n".join(block) + "\n"
            t = len(tok(text).input_ids)
            if used + t > budget_tokens:
                continue
        lines_out.extend(block)
        used += t
        if used >= budget_tokens:
            break
    return "\n".join(lines_out) if len(lines_out) > 1 else ""

def pack_context_windows(program_root: str, files: List[str], tok, budget_tokens: int) -> str:
    lines_out: List[str] = ["Program windows:"]
    used = 0
    for rel in files:
        abs_fp = rel if os.path.isabs(rel) else os.path.abspath(os.path.join(program_root, rel))
        try:
            src_lines = open(abs_fp, "r", encoding="utf-8", errors="ignore").read().splitlines()
        except Exception:
            continue
        def_lines: List[int] = []
        for i, line in enumerate(src_lines, start=1):
            s = line.lstrip()
            if s.startswith("def ") or s.startswith("class "):
                def_lines.append(i)
            if len(def_lines) >= 2:
                break
        def_lines = def_lines or [1]
        for ln in def_lines:
            a = max(1, ln - 20)
            b = min(len(src_lines), ln + 40)
            block = [f"[ctx] path: {os.path.relpath(abs_fp, program_root)}:{a}-{b}"] + src_lines[a - 1 : b] + [""]
            text = "\n".join(block) + "\n"
            t = len(tok(text).input_ids)
            if used + t > budget_tokens:
                continue
            lines_out.extend(block)
            used += t
            if used >= budget_tokens:
                break
        if used >= budget_tokens:
            break
    return "\n".join(lines_out) if len(lines_out) > 1 else ""

def collect_function_windows(program_root: str, files_: List[str], lines_each: int, *, max_candidates: int = 24) -> List[Tuple[str, int, int, int, List[str]]]:
    out: List[Tuple[str, int, int, int, List[str]]] = []
    for rel in files_:
        abs_fp = rel if os.path.isabs(rel) else os.path.abspath(os.path.join(program_root, rel))
        try:
            src_lines = open(abs_fp, "r", encoding="utf-8", errors="ignore").read().splitlines()
        except Exception:
            continue
        anchors: List[int] = []
        for i, line in enumerate(src_lines, start=1):
            s = line.lstrip()
            if s.startswith("def ") or s.startswith("class "):
                anchors.append(i)
                if len(anchors) >= max(4, int(max_candidates) // max(1, len(files_))):
                    break
        if not anchors:
            continue
        half = max(10, int(lines_each // 2))
        for ln in anchors:
            a = max(1, ln - half)
            b = min(len(src_lines), ln + half)
            out.append((rel, a, b, ln, src_lines[a - 1 : b]))
    return out

def extract_func_name_from_lines(lines_block: List[str], a: int, b: int, anchor_ln: int) -> Optional[str]:
    try:
        best_name = None
        best_dist = 10**9
        abs_ln = a
        import re as _re
        for ln_text in lines_block:
            s = ln_text.lstrip()
            if s.startswith("def ") or s.startswith("class "):
                m = _re.match(r"^(?:def|class)\s+([A-Za-z0-9_]+)", s)
                if m:
                    name = m.group(1)
                    dist = abs(abs_ln - anchor_ln)
                    if dist < best_dist:
                        best_dist = dist
                        best_name = name
            abs_ln += 1
        return best_name
    except Exception:
        return None

def score_yes_no(tok, model, q: str) -> float:
    q_ids = tok(q, return_tensors="pt")
    dev = next(model.parameters()).device
    q_ids = {k: v.to(dev) for k, v in q_ids.items()}
    with torch.no_grad():
        out_lm = model(**q_ids)
        # Support both HF-style outputs and raw tensor outputs
        if isinstance(out_lm, torch.Tensor):
            logits = out_lm
        else:
            logits = getattr(out_lm, "logits", None)
            if logits is None:
                # Best-effort fallbacks
                if isinstance(out_lm, (tuple, list)) and out_lm:
                    logits = out_lm[0]
                else:
                    logits = getattr(out_lm, "last_hidden_state", None)
            if logits is None:
                # As a last resort, return neutral 0.5
                return 0.5
        last = logits[:, -1, :]
        t1 = tok("1", add_special_tokens=False).input_ids
        t0 = tok("0", add_special_tokens=False).input_ids
        probs = torch.softmax(last, dim=-1)
        p1 = float(sum(probs[0, i].item() for i in (t1 or [])))
        p0 = float(sum(probs[0, i].item() for i in (t0 or [])))
        denom = max(1e-9, p1 + p0)
        return float(p1 / denom)

def model_prob_yes(tok, model, prompt_q: str, window_txt: str) -> Tuple[float, float]:
    rel_q = (
        "Question: " + prompt_q + "\nWindow:\n" + window_txt[:1800] + "\nDoes this window contain the core function or logic to answer the question? Answer 1 or 0."
    )
    noise_q = (
        "Question: " + prompt_q + "\nWindow:\n" + window_txt[:1800] + "\nIs this window likely test/tool/noise unrelated to answering the question? Answer 1 or 0."
    )
    return score_yes_no(tok, model, rel_q), score_yes_no(tok, model, noise_q)

""""""  # end


# function-first packer (lift from run_repo_adapter.py)

from typing import List, Optional, Tuple
import os
import torch  # type: ignore

def pack_context_heads(program_root: str, files: List[str], tok, budget_tokens: int) -> str:
    lines_out: List[str] = ["Program snippets:"]
    used = 0
    for rel in files:
        abs_fp = rel if os.path.isabs(rel) else os.path.abspath(os.path.join(program_root, rel))
        try:
            src_lines = open(abs_fp, "r", encoding="utf-8", errors="ignore").read().splitlines()
        except Exception:
            continue
        head_n = min(len(src_lines), 120)
        block = [f"[ctx] path: {os.path.relpath(abs_fp, program_root)}:1-{head_n}"] + src_lines[:head_n] + [""]
        text = "\n".join(block) + "\n"
        t = len(tok(text).input_ids)
        if used + t > budget_tokens:
            head_n = min(len(src_lines), 60)
            block = [f"[ctx] path: {os.path.relpath(abs_fp, program_root)}:1-{head_n}"] + src_lines[:head_n] + [""]
            text = "\n".join(block) + "\n"
            t = len(tok(text).input_ids)
            if used + t > budget_tokens:
                continue
        lines_out.extend(block)
        used += t
        if used >= budget_tokens:
            break
    return "\n".join(lines_out) if len(lines_out) > 1 else ""

def pack_context_windows(program_root: str, files: List[str], tok, budget_tokens: int) -> str:
    lines_out: List[str] = ["Program windows:"]
    used = 0
    for rel in files:
        abs_fp = rel if os.path.isabs(rel) else os.path.abspath(os.path.join(program_root, rel))
        try:
            src_lines = open(abs_fp, "r", encoding="utf-8", errors="ignore").read().splitlines()
        except Exception:
            continue
        def_lines: List[int] = []
        for i, line in enumerate(src_lines, start=1):
            s = line.lstrip()
            if s.startswith("def ") or s.startswith("class "):
                def_lines.append(i)
            if len(def_lines) >= 2:
                break
        def_lines = def_lines or [1]
        for ln in def_lines:


        except Exception:
            continue
        head_n = min(len(src_lines), 120)
        block = [f"[ctx] path: {os.path.relpath(abs_fp, program_root)}:1-{head_n}"] + src_lines[:head_n] + [""]
        text = "\n".join(block) + "\n"
        t = len(tok(text).input_ids)
        if used + t > budget_tokens:
            head_n = min(len(src_lines), 60)
            block = [f"[ctx] path: {os.path.relpath(abs_fp, program_root)}:1-{head_n}"] + src_lines[:head_n] + [""]
            text = "\n".join(block) + "\n"
            t = len(tok(text).input_ids)
            if used + t > budget_tokens:
                continue
        lines_out.extend(block)
        used += t
        if used >= budget_tokens:
            break
    return "\n".join(lines_out) if len(lines_out) > 1 else ""

def pack_context_windows(program_root: str, files: List[str], tok, budget_tokens: int) -> str:
    lines_out: List[str] = ["Program windows:"]
    used = 0
    for rel in files:
        abs_fp = rel if os.path.isabs(rel) else os.path.abspath(os.path.join(program_root, rel))
        try:
            src_lines = open(abs_fp, "r", encoding="utf-8", errors="ignore").read().splitlines()
        except Exception:
            continue
        def_lines: List[int] = []
        for i, line in enumerate(src_lines, start=1):
            s = line.lstrip()
            if s.startswith("def ") or s.startswith("class "):
                def_lines.append(i)
            if len(def_lines) >= 2:
                break
        def_lines = def_lines or [1]
        for ln in def_lines:
            a = max(1, ln - 20)
            b = min(len(src_lines), ln + 40)
            block = [f"[ctx] path: {os.path.relpath(abs_fp, program_root)}:{a}-{b}"] + src_lines[a - 1 : b] + [""]
            text = "\n".join(block) + "\n"
            t = len(tok(text).input_ids)
            if used + t > budget_tokens:
                continue
            lines_out.extend(block)
            used += t
            if used >= budget_tokens:
                break
        if used >= budget_tokens:
            break
    return "\n".join(lines_out) if len(lines_out) > 1 else ""

def collect_function_windows(program_root: str, files_: List[str], lines_each: int, *, max_candidates: int = 24) -> List[Tuple[str, int, int, int, List[str]]]:
    out: List[Tuple[str, int, int, int, List[str]]] = []
    for rel in files_:
        abs_fp = rel if os.path.isabs(rel) else os.path.abspath(os.path.join(program_root, rel))
        try:
            src_lines = open(abs_fp, "r", encoding="utf-8", errors="ignore").read().splitlines()
        except Exception:
            continue
        anchors: List[int] = []
        for i, line in enumerate(src_lines, start=1):
            s = line.lstrip()
            if s.startswith("def ") or s.startswith("class "):
                anchors.append(i)
                if len(anchors) >= max(4, int(max_candidates) // max(1, len(files_))):
                    break
        if not anchors:
            continue
        half = max(10, int(lines_each // 2))


            if len(def_lines) >= 2:
                break
        def_lines = def_lines or [1]
        for ln in def_lines:
            a = max(1, ln - 20)
            b = min(len(src_lines), ln + 40)
            block = [f"[ctx] path: {os.path.relpath(abs_fp, program_root)}:{a}-{b}"] + src_lines[a - 1 : b] + [""]
            text = "\n".join(block) + "\n"
            t = len(tok(text).input_ids)
            if used + t > budget_tokens:
                continue
            lines_out.extend(block)
            used += t
            if used >= budget_tokens:
                break
        if used >= budget_tokens:
            break
    return "\n".join(lines_out) if len(lines_out) > 1 else ""

def collect_function_windows(program_root: str, files_: List[str], lines_each: int, *, max_candidates: int = 24) -> List[Tuple[str, int, int, int, List[str]]]:
    out: List[Tuple[str, int, int, int, List[str]]] = []
    for rel in files_:
        abs_fp = rel if os.path.isabs(rel) else os.path.abspath(os.path.join(program_root, rel))
        try:
            src_lines = open(abs_fp, "r", encoding="utf-8", errors="ignore").read().splitlines()
        except Exception:
            continue
        anchors: List[int] = []
        for i, line in enumerate(src_lines, start=1):
            s = line.lstrip()
            if s.startswith("def ") or s.startswith("class "):
                anchors.append(i)
                if len(anchors) >= max(4, int(max_candidates) // max(1, len(files_))):
                    break
        if not anchors:
            continue
        half = max(10, int(lines_each // 2))
        for ln in anchors:
            a = max(1, ln - half)
            b = min(len(src_lines), ln + half)
            out.append((rel, a, b, ln, src_lines[a - 1 : b]))
    return out

def extract_func_name_from_lines(lines_block: List[str], a: int, b: int, anchor_ln: int) -> Optional[str]:
    try:
        best_name = None
        best_dist = 10**9
        abs_ln = a
        import re as _re
        for ln_text in lines_block:
            s = ln_text.lstrip()
            if s.startswith("def ") or s.startswith("class "):
                m = _re.match(r"^(?:def|class)\s+([A-Za-z0-9_]+)", s)
                if m:
                    name = m.group(1)
                    dist = abs(abs_ln - anchor_ln)
                    if dist < best_dist:
                        best_dist = dist
                        best_name = name
            abs_ln += 1
        return best_name


            src_lines = open(abs_fp, "r", encoding="utf-8", errors="ignore").read().splitlines()
        except Exception:
            continue
        anchors: List[int] = []
        for i, line in enumerate(src_lines, start=1):
            s = line.lstrip()
            if s.startswith("def ") or s.startswith("class "):
                anchors.append(i)
                if len(anchors) >= max(4, int(max_candidates) // max(1, len(files_))):
                    break
        if not anchors:
            continue
        half = max(10, int(lines_each // 2))
        for ln in anchors:
            a = max(1, ln - half)
            b = min(len(src_lines), ln + half)
            out.append((rel, a, b, ln, src_lines[a - 1 : b]))
    return out

def extract_func_name_from_lines(lines_block: List[str], a: int, b: int, anchor_ln: int) -> Optional[str]:
    try:
        best_name = None
        best_dist = 10**9
        abs_ln = a
        import re as _re
        for ln_text in lines_block:
            s = ln_text.lstrip()
            if s.startswith("def ") or s.startswith("class "):
                m = _re.match(r"^(?:def|class)\s+([A-Za-z0-9_]+)", s)
                if m:
                    name = m.group(1)
                    dist = abs(abs_ln - anchor_ln)
                    if dist < best_dist:
                        best_dist = dist
                        best_name = name
            abs_ln += 1
        return best_name
    except Exception:
        return None

def score_yes_no(tok, model, q: str) -> float:
    q_ids = tok(q, return_tensors="pt")
    dev = next(model.parameters()).device
    q_ids = {k: v.to(dev) for k, v in q_ids.items()}
    with torch.no_grad():
        out_lm = model(**q_ids)
        # Support both HF-style outputs and raw tensor outputs
        if isinstance(out_lm, torch.Tensor):
            logits = out_lm
        else:
            logits = getattr(out_lm, "logits", None)
            if logits is None:
                # Best-effort fallbacks
                if isinstance(out_lm, (tuple, list)) and out_lm:
                    logits = out_lm[0]
                else:
                    logits = getattr(out_lm, "last_hidden_state", None)


        best_dist = 10**9
        abs_ln = a
        import re as _re
        for ln_text in lines_block:
            s = ln_text.lstrip()
            if s.startswith("def ") or s.startswith("class "):
                m = _re.match(r"^(?:def|class)\s+([A-Za-z0-9_]+)", s)
                if m:
                    name = m.group(1)
                    dist = abs(abs_ln - anchor_ln)
                    if dist < best_dist:
                        best_dist = dist
                        best_name = name
            abs_ln += 1
        return best_name
    except Exception:
        return None

def score_yes_no(tok, model, q: str) -> float:
    q_ids = tok(q, return_tensors="pt")
    dev = next(model.parameters()).device
    q_ids = {k: v.to(dev) for k, v in q_ids.items()}
    with torch.no_grad():
        out_lm = model(**q_ids)
        # Support both HF-style outputs and raw tensor outputs
        if isinstance(out_lm, torch.Tensor):
            logits = out_lm
        else:
            logits = getattr(out_lm, "logits", None)
            if logits is None:
                # Best-effort fallbacks
                if isinstance(out_lm, (tuple, list)) and out_lm:
                    logits = out_lm[0]
                else:
                    logits = getattr(out_lm, "last_hidden_state", None)
            if logits is None:
                # As a last resort, return neutral 0.5
                return 0.5
        last = logits[:, -1, :]
        t1 = tok("1", add_special_tokens=False).input_ids
        t0 = tok("0", add_special_tokens=False).input_ids
        probs = torch.softmax(last, dim=-1)
        p1 = float(sum(probs[0, i].item() for i in (t1 or [])))
        p0 = float(sum(probs[0, i].item() for i in (t0 or [])))
        denom = max(1e-9, p1 + p0)
        return float(p1 / denom)

def model_prob_yes(tok, model, prompt_q: str, window_txt: str) -> Tuple[float, float]:
    rel_q = (
        "Question: " + prompt_q + "\nWindow:\n" + window_txt[:1800] + "\nDoes this window contain the core function or logic to answer the question? Answer 1 or 0."
    )
    noise_q = (
        "Question: " + prompt_q + "\nWindow:\n" + window_txt[:1800] + "\nIs this window likely test/tool/noise unrelated to answering the question? Answer 1 or 0."
    )
    return score_yes_no(tok, model, rel_q), score_yes_no(tok, model, noise_q)

""""""  # end


            logits = getattr(out_lm, "logits", None)
            if logits is None:
                # Best-effort fallbacks
                if isinstance(out_lm, (tuple, list)) and out_lm:
                    logits = out_lm[0]
                else:
                    logits = getattr(out_lm, "last_hidden_state", None)
            if logits is None:
                # As a last resort, return neutral 0.5
                return 0.5
        last = logits[:, -1, :]
        t1 = tok("1", add_special_tokens=False).input_ids
        t0 = tok("0", add_special_tokens=False).input_ids
        probs = torch.softmax(last, dim=-1)
        p1 = float(sum(probs[0, i].item() for i in (t1 or [])))
        p0 = float(sum(probs[0, i].item() for i in (t0 or [])))
        denom = max(1e-9, p1 + p0)
        return float(p1 / denom)

def model_prob_yes(tok, model, prompt_q: str, window_txt: str) -> Tuple[float, float]:
    rel_q = (
        "Question: " + prompt_q + "\nWindow:\n" + window_txt[:1800] + "\nDoes this window contain the core function or logic to answer the question? Answer 1 or 0."
    )
    noise_q = (
        "Question: " + prompt_q + "\nWindow:\n" + window_txt[:1800] + "\nIs this window likely test/tool/noise unrelated to answering the question? Answer 1 or 0."
    )
    return score_yes_no(tok, model, rel_q), score_yes_no(tok, model, noise_q)

""""""  # end


modules.priors

# kbann_priors derivation, round_lora utilities (move from run_repo_adapter.py/build_repo_adapter.py)


modules.targets

from typing import Optional, Dict, Tuple

def parse_target_shapes(arg: Optional[str]) -> Optional[Dict[str, Tuple[int, int]]]:
    if not arg:
        return None
    result: Dict[str, Tuple[int, int]] = {}
    try:
        parts = [p.strip() for p in str(arg).split(",") if p.strip()]
        for p in parts:
            if "=" not in p or ":" not in p:
                continue
            name, dims = p.split("=", 1)
            a, b = dims.split(":", 1)
            result[name.strip()] = (int(a), int(b))
        return result or None
    except Exception:
        return None


modules.capacity

from __future__ import annotations

from typing import Dict, List, Tuple

import numpy as np

def _parse_weights(spec: str) -> Dict[str, float]:
    # accept both "program" and "repo" as primary component key
    out: Dict[str, float] = {"program": 0.4, "subgraph": 0.4, "question": 0.2}
    try:
        parts = [p.strip() for p in str(spec).split(",") if p.strip()]
        for p in parts:
            if "=" not in p:
                continue
            k, v = p.split("=", 1)
            out[str(k).strip()] = float(v)
    except Exception:
        pass
    s = float(sum(max(0.0, v) for v in out.values()))
    if s > 0:
        out = {k: max(0.0, v) / s for k, v in out.items()}
    return out

def _safe_div(a: float, b: float) -> float:
    return (a / b) if (b > 0) else 0.0

def entropy_score(
    g: object,
    mods: List[str],
    files_rel: List[str],
    *,
    weights: str = "program=0.4,subgraph=0.4,question=0.2",
) -> Tuple[float, Dict[str, float]]:
    """Compute a 0..1 capacity score from program graph + selection.

    Components:
      - repo_component: module count (log-normalized) and import density
      - subgraph_component: selected modules' degree and breadth
      - question_component: number of files hit by the query
    """
    try:
        n_mod = float(len(getattr(g, "modules", {}) or {}))
    except Exception:
        n_mod = 0.0
    try:
        import_edges = float(sum(len(v or []) for v in getattr(g, "module_imports", {}).values()))
    except Exception:
        import_edges = 0.0

    # Program-level
    repo_mod = min(1.0, _safe_div(np.log1p(n_mod), np.log1p(1000.0)))
    repo_deg = min(1.0, _safe_div(import_edges, max(1.0, n_mod * 8.0)))
    repo_comp = float(0.5 * repo_mod + 0.5 * repo_deg)

    # Subgraph-level
    mods_list = list(mods or [])
    sub_m = float(len(mods_list))
    sub_deg_sum = 0.0
    try:
        imports = getattr(g, "module_imports", {}) or {}
    except Exception:
        imports = {}
    for m in mods_list:
        try:
            indeg = sum(1 for _x, deps in imports.items() if m in (deps or []))
            outdeg = float(len(imports.get(m, []) or []))
            sub_deg_sum += float(indeg + outdeg)
        except Exception:
            continue
    sub_deg_norm = _safe_div(sub_deg_sum, max(1.0, sub_m * 8.0))
    sub_breadth = min(1.0, _safe_div(sub_m, max(1.0, n_mod)))
    sub_comp = float(0.5 * sub_deg_norm + 0.5 * sub_breadth)

    # Question-level
    q_files = float(len(files_rel or []))
    q_comp = min(1.0, _safe_div(q_files, 24.0))

    w = _parse_weights(weights)
    prim = w.get("program", w.get("repo", 0.4))
    score = float(prim * repo_comp + w.get("subgraph", 0.4) * sub_comp + w.get("question", 0.2) * q_comp)
    diag = {
        "program_modules": n_mod,
        "program_import_edges": import_edges,
        "program_component": repo_comp,
        "subgraph_modules": sub_m,
        "subgraph_deg_norm": sub_deg_norm,
        "subgraph_component": sub_comp,
        "question_files": q_files,
        "question_component": q_comp,
        "score": score,
    }
    return max(0.0, min(1.0, score)), diag

def scale_capacity(
    es: float, *, rank_min: int, rank_max: int, gsub_min: float, gsub_max: float
) -> Tuple[int, float]:
    r = int(round(rank_min + es * max(0, (rank_max - rank_min))))
    g = float(gsub_min + es * max(0.0, (gsub_max - gsub_min)))
    r = max(rank_min, min(rank_max, r))
    g = max(min(gsub_max, max(gsub_min, g)), 0.0)
    return r, g


from __future__ import annotations

from typing import Dict, List, Tuple

import numpy as np

def _parse_weights(spec: str) -> Dict[str, float]:
    # accept both "program" and "repo" as primary component key
    out: Dict[str, float] = {"program": 0.4, "subgraph": 0.4, "question": 0.2}
    try:
        parts = [p.strip() for p in str(spec).split(",") if p.strip()]
        for p in parts:
            if "=" not in p:
                continue
            k, v = p.split("=", 1)
            out[str(k).strip()] = float(v)
    except Exception:
        pass
    s = float(sum(max(0.0, v) for v in out.values()))
    if s > 0:
        out = {k: max(0.0, v) / s for k, v in out.items()}
    return out

def _safe_div(a: float, b: float) -> float:
    return (a / b) if (b > 0) else 0.0

def entropy_score(
    g: object,
    mods: List[str],
    files_rel: List[str],
    *,
    weights: str = "program=0.4,subgraph=0.4,question=0.2",
) -> Tuple[float, Dict[str, float]]:
    """Compute a 0..1 capacity score from program graph + selection.

    Components:
      - repo_component: module count (log-normalized) and import density
      - subgraph_component: selected modules' degree and breadth
      - question_component: number of files hit by the query
    """


def _parse_weights(spec: str) -> Dict[str, float]:
    # accept both "program" and "repo" as primary component key
    out: Dict[str, float] = {"program": 0.4, "subgraph": 0.4, "question": 0.2}
    try:
        parts = [p.strip() for p in str(spec).split(",") if p.strip()]
        for p in parts:
            if "=" not in p:
                continue
            k, v = p.split("=", 1)
            out[str(k).strip()] = float(v)
    except Exception:
        pass
    s = float(sum(max(0.0, v) for v in out.values()))
    if s > 0:
        out = {k: max(0.0, v) / s for k, v in out.items()}
    return out

def _safe_div(a: float, b: float) -> float:
    return (a / b) if (b > 0) else 0.0

def entropy_score(
    g: object,
    mods: List[str],
    files_rel: List[str],
    *,
    weights: str = "program=0.4,subgraph=0.4,question=0.2",
) -> Tuple[float, Dict[str, float]]:
    """Compute a 0..1 capacity score from program graph + selection.

    Components:
      - repo_component: module count (log-normalized) and import density
      - subgraph_component: selected modules' degree and breadth
      - question_component: number of files hit by the query
    """
    try:
        n_mod = float(len(getattr(g, "modules", {}) or {}))
    except Exception:
        n_mod = 0.0


    out: Dict[str, float] = {"program": 0.4, "subgraph": 0.4, "question": 0.2}
    try:
        parts = [p.strip() for p in str(spec).split(",") if p.strip()]
        for p in parts:
            if "=" not in p:
                continue
            k, v = p.split("=", 1)
            out[str(k).strip()] = float(v)
    except Exception:
        pass
    s = float(sum(max(0.0, v) for v in out.values()))
    if s > 0:
        out = {k: max(0.0, v) / s for k, v in out.items()}
    return out

def _safe_div(a: float, b: float) -> float:
    return (a / b) if (b > 0) else 0.0

def entropy_score(
    g: object,
    mods: List[str],
    files_rel: List[str],
    *,
    weights: str = "program=0.4,subgraph=0.4,question=0.2",
) -> Tuple[float, Dict[str, float]]:
    """Compute a 0..1 capacity score from program graph + selection.

    Components:
      - repo_component: module count (log-normalized) and import density
      - subgraph_component: selected modules' degree and breadth
      - question_component: number of files hit by the query
    """
    try:
        n_mod = float(len(getattr(g, "modules", {}) or {}))
    except Exception:
        n_mod = 0.0
    try:
        import_edges = float(sum(len(v or []) for v in getattr(g, "module_imports", {}).values()))
    except Exception:
        import_edges = 0.0

    # Program-level
    repo_mod = min(1.0, _safe_div(np.log1p(n_mod), np.log1p(1000.0)))
    repo_deg = min(1.0, _safe_div(import_edges, max(1.0, n_mod * 8.0)))
    repo_comp = float(0.5 * repo_mod + 0.5 * repo_deg)

    # Subgraph-level
    mods_list = list(mods or [])
    sub_m = float(len(mods_list))
    sub_deg_sum = 0.0
    try:
        imports = getattr(g, "module_imports", {}) or {}
    except Exception:
        imports = {}
    for m in mods_list:
        try:
            indeg = sum(1 for _x, deps in imports.items() if m in (deps or []))
            outdeg = float(len(imports.get(m, []) or []))
            sub_deg_sum += float(indeg + outdeg)
        except Exception:
            continue
    sub_deg_norm = _safe_div(sub_deg_sum, max(1.0, sub_m * 8.0))
    sub_breadth = min(1.0, _safe_div(sub_m, max(1.0, n_mod)))
    sub_comp = float(0.5 * sub_deg_norm + 0.5 * sub_breadth)

    # Question-level
    q_files = float(len(files_rel or []))
    q_comp = min(1.0, _safe_div(q_files, 24.0))

    w = _parse_weights(weights)
    prim = w.get("program", w.get("repo", 0.4))
    score = float(prim * repo_comp + w.get("subgraph", 0.4) * sub_comp + w.get("question", 0.2) * q_comp)
    diag = {
        "program_modules": n_mod,
        "program_import_edges": import_edges,
        "program_component": repo_comp,
        "subgraph_modules": sub_m,
        "subgraph_deg_norm": sub_deg_norm,
        "subgraph_component": sub_comp,
        "question_files": q_files,
        "question_component": q_comp,
        "score": score,
    }
    return max(0.0, min(1.0, score)), diag

def scale_capacity(
    es: float, *, rank_min: int, rank_max: int, gsub_min: float, gsub_max: float
) -> Tuple[int, float]:
    r = int(round(rank_min + es * max(0, (rank_max - rank_min))))
    g = float(gsub_min + es * max(0.0, (gsub_max - gsub_min)))
    r = max(rank_min, min(rank_max, r))
    g = max(min(gsub_max, max(gsub_min, g)), 0.0)
    return r, g


    q_files = float(len(files_rel or []))
    q_comp = min(1.0, _safe_div(q_files, 24.0))

    w = _parse_weights(weights)
    prim = w.get("program", w.get("repo", 0.4))
    score = float(prim * repo_comp + w.get("subgraph", 0.4) * sub_comp + w.get("question", 0.2) * q_comp)
    diag = {
        "program_modules": n_mod,
        "program_import_edges": import_edges,
        "program_component": repo_comp,
        "subgraph_modules": sub_m,
        "subgraph_deg_norm": sub_deg_norm,
        "subgraph_component": sub_comp,
        "question_files": q_files,
        "question_component": q_comp,
        "score": score,
    }
    return max(0.0, min(1.0, score)), diag

def scale_capacity(
    es: float, *, rank_min: int, rank_max: int, gsub_min: float, gsub_max: float
) -> Tuple[int, float]:
    r = int(round(rank_min + es * max(0, (rank_max - rank_min))))
    g = float(gsub_min + es * max(0.0, (gsub_max - gsub_min)))
    r = max(rank_min, min(rank_max, r))
    g = max(min(gsub_max, max(gsub_min, g)), 0.0)
    return r, g


from __future__ import annotations

from dataclasses import dataclass
from typing import Dict, Tuple, Iterable, Optional, List, Set
import math

from .program_graph import ProgramGraph, Entity, Edge, Span

@dataclass
class RetrievalMix:
    sim: float = 0.6
    struct: float = 0.4
    temp: float = 0.7

    @staticmethod
    def parse(spec: Optional[str], temp: Optional[float]) -> "RetrievalMix":
        if not spec:
            return RetrievalMix(temp=(temp if isinstance(temp, (int, float)) else 0.7))
        w_sim = 0.0
        w_struct = 0.0
        try:
            parts = [p.strip() for p in str(spec).split(",") if p.strip()]
            for p in parts:
                if ":" not in p:
                    continue
                k, v = p.split(":", 1)
                if k.strip() == "sim":
                    w_sim = float(v)
                elif k.strip() == "struct":
                    w_struct = float(v)
        except Exception:
            w_sim, w_struct = 0.6, 0.4
        # normalize weights
        s = max(1e-6, (w_sim + w_struct))
        w_sim /= s
        w_struct /= s
        return RetrievalMix(sim=w_sim, struct=w_struct, temp=(temp if isinstance(temp, (int, float)) else 0.7))

class RetrievalPolicy:
    def __init__(self, mix: RetrievalMix):
        self.mix = mix

    @staticmethod
    def from_spec(policy: Optional[str], temp: Optional[float] = None) -> "RetrievalPolicy":
        return RetrievalPolicy(RetrievalMix.parse(policy, temp))

    def score_entities(self, query: str, pg: ProgramGraph) -> Dict[str, float]:
        # sim score: token overlap against entity name (very light BM25-ish)
        q_tokens = _tokenize(query)
        entities = list(pg.entities())
        sim: Dict[str, float] = {}
        for e in entities:
            name_tokens = _tokenize(e.name)
            overlap = len(q_tokens & name_tokens)
            sim[e.id] = float(overlap)
        # struct score: graph distance from top sim seeds (k=8), shorter distance => higher score
        seeds = sorted(entities, key=lambda x: sim.get(x.id, 0.0), reverse=True)[:8]
        seed_ids = [e.id for e in seeds if sim.get(e.id, 0.0) > 0]
        dist = _bfs_distance(seed_ids, pg)
        struct: Dict[str, float] = {}
        for e in entities:
            d = dist.get(e.id, None)
            if d is None:
                struct[e.id] = 0.0
            else:
                # invert distance with light decay; d=0 => 1.0
                struct[e.id] = 1.0 / float(1 + d)
        # blend and softmax with temperature
        raw: Dict[str, float] = {}
        for e in entities:
            raw[e.id] = self.mix.sim * sim.get(e.id, 0.0) + self.mix.struct * struct.get(e.id, 0.0)
        return _softmax(raw, temperature=max(1e-3, float(self.mix.temp)))

def _tokenize(text: str) -> Set[str]:
    import re
    toks = [t.lower() for t in re.findall(r"[A-Za-z0-9_]+", text or "")]
    return set(t for t in toks if len(t) > 1)

def _bfs_distance(seeds: List[str], pg: ProgramGraph, max_depth: int = 3) -> Dict[str, int]:
    # build adjacency from edges()
    adj: Dict[str, List[str]] = {}
    for e in pg.edges():
        adj.setdefault(e.src, []).append(e.dst)
        adj.setdefault(e.dst, []).append(e.src)
    dist: Dict[str, int] = {}
    cur = list(seeds)
    for s in seeds:
        dist[s] = 0
    depth = 0
    while cur and depth < max_depth:
        nxt: List[str] = []
        for u in cur:
            for v in adj.get(u, []):
                if v not in dist:
                    dist[v] = dist[u] + 1
                    nxt.append(v)
        cur = nxt
        depth += 1
    return dist

def _softmax(scores: Dict[str, float], temperature: float) -> Dict[str, float]:
    # numerically stable softmax over values
    vals = list(scores.values())
    if not vals:
        return {}
    m = max(vals)
    exps = {k: math.exp((v - m) / temperature) for k, v in scores.items()}
    s = sum(exps.values()) or 1.0
    return {k: (v / s) for k, v in exps.items()}


RetrievalMix

from __future__ import annotations

from dataclasses import dataclass
from typing import Dict, Tuple, Iterable, Optional, List, Set
import math

from .program_graph import ProgramGraph, Entity, Edge, Span

@dataclass
class RetrievalMix:
    sim: float = 0.6
    struct: float = 0.4
    temp: float = 0.7

    @staticmethod
    def parse(spec: Optional[str], temp: Optional[float]) -> "RetrievalMix":
        if not spec:
            return RetrievalMix(temp=(temp if isinstance(temp, (int, float)) else 0.7))
        w_sim = 0.0
        w_struct = 0.0
        try:
            parts = [p.strip() for p in str(spec).split(",") if p.strip()]
            for p in parts:
                if ":" not in p:
                    continue
                k, v = p.split(":", 1)
                if k.strip() == "sim":
                    w_sim = float(v)
                elif k.strip() == "struct":
                    w_struct = float(v)
        except Exception:
            w_sim, w_struct = 0.6, 0.4
        # normalize weights
        s = max(1e-6, (w_sim + w_struct))
        w_sim /= s
        w_struct /= s
        return RetrievalMix(sim=w_sim, struct=w_struct, temp=(temp if isinstance(temp, (int, float)) else 0.7))

class RetrievalPolicy:
    def __init__(self, mix: RetrievalMix):
        self.mix = mix

    @staticmethod
    def from_spec(policy: Optional[str], temp: Optional[float] = None) -> "RetrievalPolicy":
        return RetrievalPolicy(RetrievalMix.parse(policy, temp))

    def score_entities(self, query: str, pg: ProgramGraph) -> Dict[str, float]:
        # sim score: token overlap against entity name (very light BM25-ish)
        q_tokens = _tokenize(query)
        entities = list(pg.entities())
        sim: Dict[str, float] = {}
        for e in entities:
            name_tokens = _tokenize(e.name)
            overlap = len(q_tokens & name_tokens)
            sim[e.id] = float(overlap)
        # struct score: graph distance from top sim seeds (k=8), shorter distance => higher score


RetrievalPolicy

        w_struct = 0.0
        try:
            parts = [p.strip() for p in str(spec).split(",") if p.strip()]
            for p in parts:
                if ":" not in p:
                    continue
                k, v = p.split(":", 1)
                if k.strip() == "sim":
                    w_sim = float(v)
                elif k.strip() == "struct":
                    w_struct = float(v)
        except Exception:
            w_sim, w_struct = 0.6, 0.4
        # normalize weights
        s = max(1e-6, (w_sim + w_struct))
        w_sim /= s
        w_struct /= s
        return RetrievalMix(sim=w_sim, struct=w_struct, temp=(temp if isinstance(temp, (int, float)) else 0.7))

class RetrievalPolicy:
    def __init__(self, mix: RetrievalMix):
        self.mix = mix

    @staticmethod
    def from_spec(policy: Optional[str], temp: Optional[float] = None) -> "RetrievalPolicy":
        return RetrievalPolicy(RetrievalMix.parse(policy, temp))

    def score_entities(self, query: str, pg: ProgramGraph) -> Dict[str, float]:
        # sim score: token overlap against entity name (very light BM25-ish)
        q_tokens = _tokenize(query)
        entities = list(pg.entities())
        sim: Dict[str, float] = {}
        for e in entities:
            name_tokens = _tokenize(e.name)
            overlap = len(q_tokens & name_tokens)
            sim[e.id] = float(overlap)
        # struct score: graph distance from top sim seeds (k=8), shorter distance => higher score
        seeds = sorted(entities, key=lambda x: sim.get(x.id, 0.0), reverse=True)[:8]
        seed_ids = [e.id for e in seeds if sim.get(e.id, 0.0) > 0]
        dist = _bfs_distance(seed_ids, pg)
        struct: Dict[str, float] = {}
        for e in entities:
            d = dist.get(e.id, None)
            if d is None:
                struct[e.id] = 0.0
            else:
                # invert distance with light decay; d=0 => 1.0
                struct[e.id] = 1.0 / float(1 + d)
        # blend and softmax with temperature
        raw: Dict[str, float] = {}
        for e in entities:
            raw[e.id] = self.mix.sim * sim.get(e.id, 0.0) + self.mix.struct * struct.get(e.id, 0.0)
        return _softmax(raw, temperature=max(1e-3, float(self.mix.temp)))

def _tokenize(text: str) -> Set[str]:
    import re
    toks = [t.lower() for t in re.findall(r"[A-Za-z0-9_]+", text or "")]
    return set(t for t in toks if len(t) > 1)

def _bfs_distance(seeds: List[str], pg: ProgramGraph, max_depth: int = 3) -> Dict[str, int]:
    # build adjacency from edges()
    adj: Dict[str, List[str]] = {}
    for e in pg.edges():
        adj.setdefault(e.src, []).append(e.dst)
        adj.setdefault(e.dst, []).append(e.src)
    dist: Dict[str, int] = {}
    cur = list(seeds)
    for s in seeds:
        dist[s] = 0
    depth = 0
    while cur and depth < max_depth:


            sim[e.id] = float(overlap)
        # struct score: graph distance from top sim seeds (k=8), shorter distance => higher score
        seeds = sorted(entities, key=lambda x: sim.get(x.id, 0.0), reverse=True)[:8]
        seed_ids = [e.id for e in seeds if sim.get(e.id, 0.0) > 0]
        dist = _bfs_distance(seed_ids, pg)
        struct: Dict[str, float] = {}
        for e in entities:
            d = dist.get(e.id, None)
            if d is None:
                struct[e.id] = 0.0
            else:
                # invert distance with light decay; d=0 => 1.0
                struct[e.id] = 1.0 / float(1 + d)
        # blend and softmax with temperature
        raw: Dict[str, float] = {}
        for e in entities:
            raw[e.id] = self.mix.sim * sim.get(e.id, 0.0) + self.mix.struct * struct.get(e.id, 0.0)
        return _softmax(raw, temperature=max(1e-3, float(self.mix.temp)))

def _tokenize(text: str) -> Set[str]:
    import re
    toks = [t.lower() for t in re.findall(r"[A-Za-z0-9_]+", text or "")]
    return set(t for t in toks if len(t) > 1)

def _bfs_distance(seeds: List[str], pg: ProgramGraph, max_depth: int = 3) -> Dict[str, int]:
    # build adjacency from edges()
    adj: Dict[str, List[str]] = {}
    for e in pg.edges():
        adj.setdefault(e.src, []).append(e.dst)
        adj.setdefault(e.dst, []).append(e.src)
    dist: Dict[str, int] = {}
    cur = list(seeds)
    for s in seeds:
        dist[s] = 0
    depth = 0
    while cur and depth < max_depth:
        nxt: List[str] = []
        for u in cur:
            for v in adj.get(u, []):
                if v not in dist:
                    dist[v] = dist[u] + 1
                    nxt.append(v)


        for e in entities:
            d = dist.get(e.id, None)
            if d is None:
                struct[e.id] = 0.0
            else:
                # invert distance with light decay; d=0 => 1.0
                struct[e.id] = 1.0 / float(1 + d)
        # blend and softmax with temperature
        raw: Dict[str, float] = {}
        for e in entities:
            raw[e.id] = self.mix.sim * sim.get(e.id, 0.0) + self.mix.struct * struct.get(e.id, 0.0)
        return _softmax(raw, temperature=max(1e-3, float(self.mix.temp)))

def _tokenize(text: str) -> Set[str]:
    import re
    toks = [t.lower() for t in re.findall(r"[A-Za-z0-9_]+", text or "")]
    return set(t for t in toks if len(t) > 1)

def _bfs_distance(seeds: List[str], pg: ProgramGraph, max_depth: int = 3) -> Dict[str, int]:
    # build adjacency from edges()
    adj: Dict[str, List[str]] = {}
    for e in pg.edges():
        adj.setdefault(e.src, []).append(e.dst)
        adj.setdefault(e.dst, []).append(e.src)
    dist: Dict[str, int] = {}
    cur = list(seeds)
    for s in seeds:
        dist[s] = 0
    depth = 0
    while cur and depth < max_depth:
        nxt: List[str] = []
        for u in cur:
            for v in adj.get(u, []):
                if v not in dist:
                    dist[v] = dist[u] + 1
                    nxt.append(v)
        cur = nxt
        depth += 1
    return dist

def _softmax(scores: Dict[str, float], temperature: float) -> Dict[str, float]:
    # numerically stable softmax over values
    vals = list(scores.values())
    if not vals:
        return {}
    m = max(vals)
    exps = {k: math.exp((v - m) / temperature) for k, v in scores.items()}
    s = sum(exps.values()) or 1.0
    return {k: (v / s) for k, v in exps.items()}


    for e in pg.edges():
        adj.setdefault(e.src, []).append(e.dst)
        adj.setdefault(e.dst, []).append(e.src)
    dist: Dict[str, int] = {}
    cur = list(seeds)
    for s in seeds:
        dist[s] = 0
    depth = 0
    while cur and depth < max_depth:
        nxt: List[str] = []
        for u in cur:
            for v in adj.get(u, []):
                if v not in dist:
                    dist[v] = dist[u] + 1
                    nxt.append(v)
        cur = nxt
        depth += 1
    return dist

def _softmax(scores: Dict[str, float], temperature: float) -> Dict[str, float]:
    # numerically stable softmax over values
    vals = list(scores.values())
    if not vals:
        return {}
    m = max(vals)
    exps = {k: math.exp((v - m) / temperature) for k, v in scores.items()}
    s = sum(exps.values()) or 1.0
    return {k: (v / s) for k, v in exps.items()}


        try:
            parts = [p.strip() for p in str(spec).split(",") if p.strip()]
            for p in parts:
                if ":" not in p:
                    continue
                k, v = p.split(":", 1)
                if k.strip() == "sim":
                    w_sim = float(v)
                elif k.strip() == "struct":
                    w_struct = float(v)
        except Exception:
            w_sim, w_struct = 0.6, 0.4
        # normalize weights
        s = max(1e-6, (w_sim + w_struct))
        w_sim /= s
        w_struct /= s
        return RetrievalMix(sim=w_sim, struct=w_struct, temp=(temp if isinstance(temp, (int, float)) else 0.7))

class RetrievalPolicy:
    def __init__(self, mix: RetrievalMix):
        self.mix = mix

    @staticmethod
    def from_spec(policy: Optional[str], temp: Optional[float] = None) -> "RetrievalPolicy":
        return RetrievalPolicy(RetrievalMix.parse(policy, temp))

    def score_entities(self, query: str, pg: ProgramGraph) -> Dict[str, float]:
        # sim score: token overlap against entity name (very light BM25-ish)
        q_tokens = _tokenize(query)
        entities = list(pg.entities())
        sim: Dict[str, float] = {}
        for e in entities:
            name_tokens = _tokenize(e.name)
            overlap = len(q_tokens & name_tokens)
            sim[e.id] = float(overlap)
        # struct score: graph distance from top sim seeds (k=8), shorter distance => higher score
        seeds = sorted(entities, key=lambda x: sim.get(x.id, 0.0), reverse=True)[:8]
        seed_ids = [e.id for e in seeds if sim.get(e.id, 0.0) > 0]
        dist = _bfs_distance(seed_ids, pg)
        struct: Dict[str, float] = {}
        for e in entities:


                    continue
                k, v = p.split(":", 1)
                if k.strip() == "sim":
                    w_sim = float(v)
                elif k.strip() == "struct":
                    w_struct = float(v)
        except Exception:
            w_sim, w_struct = 0.6, 0.4
        # normalize weights
        s = max(1e-6, (w_sim + w_struct))
        w_sim /= s
        w_struct /= s
        return RetrievalMix(sim=w_sim, struct=w_struct, temp=(temp if isinstance(temp, (int, float)) else 0.7))

class RetrievalPolicy:
    def __init__(self, mix: RetrievalMix):
        self.mix = mix

    @staticmethod
    def from_spec(policy: Optional[str], temp: Optional[float] = None) -> "RetrievalPolicy":
        return RetrievalPolicy(RetrievalMix.parse(policy, temp))

    def score_entities(self, query: str, pg: ProgramGraph) -> Dict[str, float]:
        # sim score: token overlap against entity name (very light BM25-ish)
        q_tokens = _tokenize(query)
        entities = list(pg.entities())
        sim: Dict[str, float] = {}
        for e in entities:
            name_tokens = _tokenize(e.name)
            overlap = len(q_tokens & name_tokens)
            sim[e.id] = float(overlap)
        # struct score: graph distance from top sim seeds (k=8), shorter distance => higher score
        seeds = sorted(entities, key=lambda x: sim.get(x.id, 0.0), reverse=True)[:8]
        seed_ids = [e.id for e in seeds if sim.get(e.id, 0.0) > 0]
        dist = _bfs_distance(seed_ids, pg)
        struct: Dict[str, float] = {}
        for e in entities:
            d = dist.get(e.id, None)
            if d is None:
                struct[e.id] = 0.0
            else:


                    w_sim = float(v)
                elif k.strip() == "struct":
                    w_struct = float(v)
        except Exception:
            w_sim, w_struct = 0.6, 0.4
        # normalize weights
        s = max(1e-6, (w_sim + w_struct))
        w_sim /= s
        w_struct /= s
        return RetrievalMix(sim=w_sim, struct=w_struct, temp=(temp if isinstance(temp, (int, float)) else 0.7))

class RetrievalPolicy:
    def __init__(self, mix: RetrievalMix):
        self.mix = mix

    @staticmethod
    def from_spec(policy: Optional[str], temp: Optional[float] = None) -> "RetrievalPolicy":
        return RetrievalPolicy(RetrievalMix.parse(policy, temp))

    def score_entities(self, query: str, pg: ProgramGraph) -> Dict[str, float]:
        # sim score: token overlap against entity name (very light BM25-ish)
        q_tokens = _tokenize(query)
        entities = list(pg.entities())
        sim: Dict[str, float] = {}
        for e in entities:
            name_tokens = _tokenize(e.name)
            overlap = len(q_tokens & name_tokens)
            sim[e.id] = float(overlap)
        # struct score: graph distance from top sim seeds (k=8), shorter distance => higher score
        seeds = sorted(entities, key=lambda x: sim.get(x.id, 0.0), reverse=True)[:8]
        seed_ids = [e.id for e in seeds if sim.get(e.id, 0.0) > 0]
        dist = _bfs_distance(seed_ids, pg)
        struct: Dict[str, float] = {}
        for e in entities:
            d = dist.get(e.id, None)
            if d is None:
                struct[e.id] = 0.0
            else:
                # invert distance with light decay; d=0 => 1.0
                struct[e.id] = 1.0 / float(1 + d)
        # blend and softmax with temperature
        raw: Dict[str, float] = {}
        for e in entities:
            raw[e.id] = self.mix.sim * sim.get(e.id, 0.0) + self.mix.struct * struct.get(e.id, 0.0)
        return _softmax(raw, temperature=max(1e-3, float(self.mix.temp)))

def _tokenize(text: str) -> Set[str]:
    import re
    toks = [t.lower() for t in re.findall(r"[A-Za-z0-9_]+", text or "")]
    return set(t for t in toks if len(t) > 1)

def _bfs_distance(seeds: List[str], pg: ProgramGraph, max_depth: int = 3) -> Dict[str, int]:
    # build adjacency from edges()
    adj: Dict[str, List[str]] = {}
    for e in pg.edges():
        adj.setdefault(e.src, []).append(e.dst)
        adj.setdefault(e.dst, []).append(e.src)
    dist: Dict[str, int] = {}
    cur = list(seeds)
    for s in seeds:
        dist[s] = 0
    depth = 0
    while cur and depth < max_depth:


modules.provenance

import subprocess
from typing import Optional

def git_commit_sha(repo_root: str) -> Optional[str]:
    try:
        sha = subprocess.check_output(
            ["git", "-C", repo_root, "rev-parse", "HEAD"], stderr=subprocess.DEVNULL, text=True
        ).strip()
        return sha if sha else None
    except Exception:
        return None

def git_tree_sha(repo_root: str) -> Optional[str]:
    """Return the HEAD tree SHA if available (pins exact tracked file set)."""
    try:
        sha = subprocess.check_output(
            ["git", "-C", repo_root, "rev-parse", "HEAD^{tree}"], stderr=subprocess.DEVNULL, text=True
        ).strip()
        return sha if sha else None
    except Exception:
        return None


modules.adapter

from typing import Any, Dict, List, Optional, Tuple
import numpy as np
import math
import os
import json
import re
from .embedding import _stable_hash

def _make_random_matrix(shape: Tuple[int, int], *, seed: int) -> np.ndarray:
    rng = np.random.default_rng(seed)
    # Xavier uniform
    limit = math.sqrt(6.0 / float(shape[0] + shape[1]))
    return rng.uniform(-limit, limit, size=shape).astype(np.float32)

def generate_lora_from_embedding(
    z: np.ndarray,
    *,
    d_model: int,
    num_layers: int,
    rank: int = 8,
    seed: int = 0,
    targets: Optional[List[str]] = None,
    target_shapes: Optional[Dict[str, Tuple[int, int]]] = None,
    layer_gate: str = "zmean",
    target_weights: Optional[Dict[str, float]] = None,
    learn_bias: bool = False,
    map_cap: Optional[float] = None,
) -> Dict[str, List[Dict[str, Dict[str, np.ndarray]]]]:
    if targets is None:
        targets = ["q_proj", "o_proj", "up_proj"]

    z = z.astype(np.float32)
    # Always-on normalization for mapping
    try:
        nz = float(np.linalg.norm(z))
        if nz > 0:
            z = (z / nz).astype(np.float32)
    except Exception:
        pass
    gates: List[float] = []
    layers: List[Dict[str, Dict[str, np.ndarray]]] = []

    for layer_idx in range(num_layers):
        layer_state: Dict[str, Dict[str, np.ndarray]] = {}
        # derive per-layer seeds from z by hashing a projection
        key = int((_stable_hash(f"layer:{layer_idx}", seed) ^ _stable_hash(str(float(z[0])), seed + 7)) & ((1 << 31) - 1))
        # gate schedule
        frac = float(layer_idx) / float(max(1, num_layers - 1))
        if layer_gate == "cosine":
            gate = float(0.5 * (1.0 - math.cos(math.pi * frac)))
        elif layer_gate == "hump":
            gate = float(max(0.0, math.sin(math.pi * frac)))
        elif layer_gate == "linear":
            gate = float(frac)
        else:  # zmean
            gate = float((np.tanh(z[(layer_idx * 13) % len(z)]) + 1.0) * 0.5)
        gates.append(gate)
        # Pair MLP projections: reuse seed and gate across up/gate/down; up/gate share A/B
        mlp_seed = key ^ _stable_hash("mlp_pair", seed)
        up_pair: Optional[Tuple[np.ndarray, np.ndarray]] = None
        # Precompute per-target scalar from disjoint z segments (always-on)
        per_target_scale: Dict[str, float] = {}
        try:
            seg_len = max(8, int(len(z) // max(1, len(targets))))
            for idx_t, tgt in enumerate(targets):
                start = (layer_idx * 17 + idx_t * seg_len) % len(z)
                idxs = (np.arange(seg_len) + start) % len(z)
                seg = z[idxs]
                per_target_scale[tgt] = float(np.tanh(np.mean(seg) * 1.5))
        except Exception:
            per_target_scale = {}
        for tgt in targets:
            # A: d_out x r ; B: r x d_in
            if target_shapes and tgt in target_shapes:
                d_out, d_in = target_shapes[tgt]
            else:
                d_out, d_in = d_model, d_model
            # Coupled seeds for MLP blocks to align up/gate/down
            if tgt in ("up_proj", "gate_proj"):
                if up_pair is None:
                    A = _make_random_matrix((int(d_out), rank), seed=mlp_seed ^ _stable_hash("up:A", seed))
                    B = _make_random_matrix((rank, int(d_in)), seed=mlp_seed ^ _stable_hash("up:B", seed + 1))
                    up_pair = (A, B)
                else:
                    A, B = up_pair
            elif tgt == "down_proj":
                A = _make_random_matrix((int(d_out), rank), seed=mlp_seed ^ _stable_hash("down:A", seed))
                B = _make_random_matrix((rank, int(d_in)), seed=mlp_seed ^ _stable_hash("down:B", seed + 1))
            else:
                A = _make_random_matrix((int(d_out), rank), seed=key ^ _stable_hash(tgt + ":A", seed))
                B = _make_random_matrix((rank, int(d_in)), seed=key ^ _stable_hash(tgt + ":B", seed + 1))
            # Fan-in/fan-out scaling: A *= 1/sqrt(rank); optional B zeroing controlled by zero_B flag at call site via target_weights special key
            if rank > 0:
                A = (1.0 / float(max(1.0, math.sqrt(float(rank))))) * A
            # modulate by low-d projection of z (wrap-safe segment)
            start = (layer_idx * 31) % len(z)
            idx = (np.arange(32) + start) % len(z)
            seg = z[idx]
            alpha = float(np.clip(np.mean(seg) * 1.5, -1.0, 1.0))
            A = (1.0 + alpha * gate) * A
            B = (1.0 - alpha * gate) * B
            # Enhanced: per-target scalar (light) to reduce interference across targets
            if per_target_scale and tgt in per_target_scale:
                s_t = float(per_target_scale[tgt])
                A = (1.0 + 0.15 * s_t) * A
                B = (1.0 - 0.15 * s_t) * B
            if target_weights and tgt in target_weights:
                tw = float(target_weights[tgt])
                s = float(max(0.0, tw)) ** 0.5
                A = s * A
                B = s * B
            # Optional norm cap on A/B (Frobenius) to avoid runaway scales
            try:
                if map_cap is not None and float(map_cap) > 0:
                    capv = float(map_cap)
                    nA = float(np.linalg.norm(A)) if A.size > 0 else 0.0
                    if nA > capv and nA > 0:
                        A = (capv / nA) * A
                    nB = float(np.linalg.norm(B)) if B.size > 0 else 0.0
                    if nB > capv and nB > 0:
                        B = (capv / nB) * B
            except Exception:
                pass
            e: Dict[str, np.ndarray] = {"A": A, "B": B, "gate": np.array([gate], dtype=np.float32)}
            if learn_bias:
                e["bias"] = np.zeros((int(d_out),), dtype=np.float32)
            layer_state[tgt] = e
        layers.append(layer_state)

    return {"layers": layers, "rank": rank, "d_model": d_model, "targets": targets, "gates": np.array(gates, dtype=np.float32)}

def generate_lora_from_embedding_torch(
    z: np.ndarray,
    *,
    d_model: int,
    num_layers: int,
    rank: int = 8,
    seed: int = 0,
    targets: Optional[List[str]] = None,
    target_shapes: Optional[Dict[str, Tuple[int, int]]] = None,
    einsum_opt: str = "auto",
    layer_gate: str = "zmean",
    target_weights: Optional[Dict[str, float]] = None,
    map_cap: Optional[float] = None,
) -> Dict[str, List[Dict[str, Dict[str, np.ndarray]]]]:
    import torch  # local import to avoid hard dep when unused

    if targets is None:
        targets = ["q_proj", "o_proj", "up_proj"]

    zt = torch.from_numpy(z.astype(np.float32))
    try:
        nz = float(torch.linalg.vector_norm(zt).item())
        if nz > 0:
            zt = (zt / nz).to(torch.float32)
    except Exception:
        pass
    layers: List[Dict[str, Dict[str, np.ndarray]]] = []
    gates: List[float] = []
    # Global seeding for reproducibility in case external torch ops run
    try:
        torch.manual_seed(int(seed))
    except Exception:
        pass

    for layer_idx in range(num_layers):
        layer_state: Dict[str, Dict[str, np.ndarray]] = {}
        key = int((_stable_hash(f"layer:{layer_idx}", seed) ^ _stable_hash(str(float(z[0])), seed + 7)) & ((1 << 31) - 1))
        frac = float(layer_idx) / float(max(1, num_layers - 1))
        if layer_gate == "cosine":
            gate = float(0.5 * (1.0 - math.cos(math.pi * frac)))
        elif layer_gate == "hump":
            gate = float(max(0.0, math.sin(math.pi * frac)))
        elif layer_gate == "linear":
            gate = float(frac)
        else:
            gate = float((np.tanh(z[(layer_idx * 13) % len(z)]) + 1.0) * 0.5)
        gates.append(gate)
        # deterministic torch RNG
        gen = torch.Generator(device="cpu")
        gen.manual_seed(key)
        mlp_seed = key ^ _stable_hash("mlp_pair", seed)
        up_pair: Optional[Tuple[torch.Tensor, torch.Tensor]] = None
        # Precompute per-target scale (always-on)
        per_target_scale: Dict[str, float] = {}
        try:
            seg_len = max(8, int(len(zt) // max(1, len(targets))))
            for idx_t, tgt in enumerate(targets):
                start = (layer_idx * 17 + idx_t * seg_len) % len(zt)
                idxs = (torch.arange(seg_len) + start) % len(zt)
                seg = zt[idxs]
                per_target_scale[tgt] = float(torch.tanh(seg.mean() * 1.5).item())
        except Exception:
            per_target_scale = {}
        for tgt in targets:
            if target_shapes and tgt in target_shapes:
                d_out, d_in = target_shapes[tgt]
            else:
                d_out, d_in = d_model, d_model
            if tgt in ("up_proj", "gate_proj"):
                if up_pair is None:
                    A = (torch.rand((int(d_out), rank), generator=gen) * 2 - 1).to(torch.float32)
                    B = (torch.rand((rank, int(d_in)), generator=gen) * 2 - 1).to(torch.float32)
                    up_pair = (A, B)
                else:
                    A, B = up_pair
            elif tgt == "down_proj":
                # reuse generator but distinct suffix for down
                A = (torch.rand((int(d_out), rank), generator=gen) * 2 - 1).to(torch.float32)
                B = (torch.rand((rank, int(d_in)), generator=gen) * 2 - 1).to(torch.float32)
            else:
                A = (torch.rand((int(d_out), rank), generator=gen) * 2 - 1).to(torch.float32)
                B = (torch.rand((rank, int(d_in)), generator=gen) * 2 - 1).to(torch.float32)
            if rank > 0:
                A = A * (1.0 / float(max(1.0, math.sqrt(float(rank)))))
            # Modulate by contraction with a deterministic kernel vector derived from z
            start = (layer_idx * 31) % len(z)
            idx = (np.arange(32) + start) % len(z)
            seg = zt[idx]
            w = torch.sin(torch.linspace(0, 3.14159, steps=32))
            if _contract is not None and einsum_opt:
                alpha = torch.tanh(_contract("i,i->", seg, w, optimize=einsum_opt) / 8.0)
            else:
                alpha = torch.tanh((seg * w).sum() / 8.0)
            A = (1.0 + float(alpha.item()) * gate) * A
            B = (1.0 - float(alpha.item()) * gate) * B
            if per_target_scale and tgt in per_target_scale:
                s_t = float(per_target_scale[tgt])
                A = (1.0 + 0.15 * s_t) * A
                B = (1.0 - 0.15 * s_t) * B
            if target_weights and tgt in target_weights:
                tw = float(target_weights[tgt])
                s = float(max(0.0, tw)) ** 0.5
                A = (s * A)
                B = (s * B)
            # Optional norm cap (Frobenius) on A/B
            try:
                if map_cap is not None and float(map_cap) > 0:
                    capv = float(map_cap)
                    nA = float(torch.linalg.matrix_norm(A).item()) if A.numel() > 0 else 0.0
                    if nA > capv and nA > 0:
                        A = (capv / nA) * A
                    nB = float(torch.linalg.matrix_norm(B).item()) if B.numel() > 0 else 0.0
                    if nB > capv and nB > 0:
                        B = (capv / nB) * B
            except Exception:
                pass
            e: Dict[str, np.ndarray] = {
                "A": A.numpy().astype(np.float32),
                "B": B.numpy().astype(np.float32),
                "gate": np.array([gate], dtype=np.float32),
            }
            # Torch path does not currently support learn_bias flag; add zero bias for parity if requested via target_weights special key later if needed
            layer_state[tgt] = e
        layers.append(layer_state)

    return {"layers": layers, "rank": rank, "d_model": d_model, "targets": targets, "gates": np.array(gates, dtype=np.float32)}

def save_npz(out_dir: str, *, embedding: Dict[str, np.ndarray], adapters: Dict[str, Any], manifest: Dict[str, Any]) -> None:
    os.makedirs(out_dir, exist_ok=True)
    np.savez_compressed(os.path.join(out_dir, "embedding.npz"), **embedding)
    # flatten adapter arrays
    flat: Dict[str, np.ndarray] = {}
    for i, layer in enumerate(adapters["layers"]):
        for name, tensors in layer.items():
            flat[f"L{i}.{name}.A"] = tensors["A"]
            flat[f"L{i}.{name}.B"] = tensors["B"]
            gate_val = float(tensors["gate"][0]) if isinstance(tensors.get("gate"), np.ndarray) else float(tensors.get("gate", 0.0))
            flat[f"L{i}.{name}.gate"] = np.array(gate_val, dtype=np.float32)
            if "bias" in tensors and isinstance(tensors["bias"], np.ndarray):
                flat[f"L{i}.{name}.bias"] = tensors["bias"]
    np.savez_compressed(os.path.join(out_dir, "adapters.npz"), **flat)
    open(os.path.join(out_dir, "manifest.json"), "w", encoding="utf-8").write(json.dumps(manifest, indent=2))

def load_adapters_npz(path: str) -> Dict[str, List[Dict[str, Dict[str, np.ndarray]]]]:
    data = np.load(path)
    # infer indices
    layers: Dict[int, Dict[str, Dict[str, np.ndarray]]] = {}
    for key in data.files:
        # L{idx}.{name}.{A|B|gate}
        parts = key.split(".")
        if len(parts) != 3:
            continue
        lid = int(parts[0][1:])
        name = parts[1]
        kind = parts[2]
        layers.setdefault(lid, {}).setdefault(name, {})[kind] = data[key]
    ordered = [layers[i] for i in sorted(layers.keys())]
    return {"layers": ordered}


from typing import Any, Dict, List, Optional, Tuple
import numpy as np
import math
import os
import json
import re
from .embedding import _stable_hash

def _make_random_matrix(shape: Tuple[int, int], *, seed: int) -> np.ndarray:
    rng = np.random.default_rng(seed)
    # Xavier uniform
    limit = math.sqrt(6.0 / float(shape[0] + shape[1]))
    return rng.uniform(-limit, limit, size=shape).astype(np.float32)

def generate_lora_from_embedding(
    z: np.ndarray,
    *,
    d_model: int,
    num_layers: int,
    rank: int = 8,
    seed: int = 0,
    targets: Optional[List[str]] = None,
    target_shapes: Optional[Dict[str, Tuple[int, int]]] = None,
    layer_gate: str = "zmean",
    target_weights: Optional[Dict[str, float]] = None,
    learn_bias: bool = False,
    map_cap: Optional[float] = None,
) -> Dict[str, List[Dict[str, Dict[str, np.ndarray]]]]:
    if targets is None:
        targets = ["q_proj", "o_proj", "up_proj"]

    z = z.astype(np.float32)
    # Always-on normalization for mapping


from typing import Any, Dict, List, Optional, Tuple
import numpy as np
import math
import os
import json
import re
from .embedding import _stable_hash

def _make_random_matrix(shape: Tuple[int, int], *, seed: int) -> np.ndarray:
    rng = np.random.default_rng(seed)
    # Xavier uniform
    limit = math.sqrt(6.0 / float(shape[0] + shape[1]))
    return rng.uniform(-limit, limit, size=shape).astype(np.float32)

def generate_lora_from_embedding(
    z: np.ndarray,
    *,
    d_model: int,
    num_layers: int,
    rank: int = 8,
    seed: int = 0,
    targets: Optional[List[str]] = None,
    target_shapes: Optional[Dict[str, Tuple[int, int]]] = None,
    layer_gate: str = "zmean",
    target_weights: Optional[Dict[str, float]] = None,
    learn_bias: bool = False,
    map_cap: Optional[float] = None,
) -> Dict[str, List[Dict[str, Dict[str, np.ndarray]]]]:
    if targets is None:
        targets = ["q_proj", "o_proj", "up_proj"]

    z = z.astype(np.float32)
    # Always-on normalization for mapping
    try:
        nz = float(np.linalg.norm(z))
        if nz > 0:
            z = (z / nz).astype(np.float32)
    except Exception:
        pass
    gates: List[float] = []
    layers: List[Dict[str, Dict[str, np.ndarray]]] = []

    for layer_idx in range(num_layers):
        layer_state: Dict[str, Dict[str, np.ndarray]] = {}
        # derive per-layer seeds from z by hashing a projection
        key = int((_stable_hash(f"layer:{layer_idx}", seed) ^ _stable_hash(str(float(z[0])), seed + 7)) & ((1 << 31) - 1))
        # gate schedule
        frac = float(layer_idx) / float(max(1, num_layers - 1))
        if layer_gate == "cosine":
            gate = float(0.5 * (1.0 - math.cos(math.pi * frac)))
        elif layer_gate == "hump":
            gate = float(max(0.0, math.sin(math.pi * frac)))
        elif layer_gate == "linear":
            gate = float(frac)
        else:  # zmean
            gate = float((np.tanh(z[(layer_idx * 13) % len(z)]) + 1.0) * 0.5)
        gates.append(gate)
        # Pair MLP projections: reuse seed and gate across up/gate/down; up/gate share A/B
        mlp_seed = key ^ _stable_hash("mlp_pair", seed)
        up_pair: Optional[Tuple[np.ndarray, np.ndarray]] = None
        # Precompute per-target scalar from disjoint z segments (always-on)
        per_target_scale: Dict[str, float] = {}
        try:
            seg_len = max(8, int(len(z) // max(1, len(targets))))
            for idx_t, tgt in enumerate(targets):
                start = (layer_idx * 17 + idx_t * seg_len) % len(z)
                idxs = (np.arange(seg_len) + start) % len(z)
                seg = z[idxs]
                per_target_scale[tgt] = float(np.tanh(np.mean(seg) * 1.5))
        except Exception:
            per_target_scale = {}
        for tgt in targets:
            # A: d_out x r ; B: r x d_in
            if target_shapes and tgt in target_shapes:
                d_out, d_in = target_shapes[tgt]
            else:
                d_out, d_in = d_model, d_model
            # Coupled seeds for MLP blocks to align up/gate/down
            if tgt in ("up_proj", "gate_proj"):
                if up_pair is None:
                    A = _make_random_matrix((int(d_out), rank), seed=mlp_seed ^ _stable_hash("up:A", seed))
                    B = _make_random_matrix((rank, int(d_in)), seed=mlp_seed ^ _stable_hash("up:B", seed + 1))
                    up_pair = (A, B)
                else:
                    A, B = up_pair
            elif tgt == "down_proj":
                A = _make_random_matrix((int(d_out), rank), seed=mlp_seed ^ _stable_hash("down:A", seed))
                B = _make_random_matrix((rank, int(d_in)), seed=mlp_seed ^ _stable_hash("down:B", seed + 1))
            else:
                A = _make_random_matrix((int(d_out), rank), seed=key ^ _stable_hash(tgt + ":A", seed))
                B = _make_random_matrix((rank, int(d_in)), seed=key ^ _stable_hash(tgt + ":B", seed + 1))
            # Fan-in/fan-out scaling: A *= 1/sqrt(rank); optional B zeroing controlled by zero_B flag at call site via target_weights special key
            if rank > 0:
                A = (1.0 / float(max(1.0, math.sqrt(float(rank))))) * A
            # modulate by low-d projection of z (wrap-safe segment)
            start = (layer_idx * 31) % len(z)
            idx = (np.arange(32) + start) % len(z)
            seg = z[idx]
            alpha = float(np.clip(np.mean(seg) * 1.5, -1.0, 1.0))
            A = (1.0 + alpha * gate) * A
            B = (1.0 - alpha * gate) * B
            # Enhanced: per-target scalar (light) to reduce interference across targets
            if per_target_scale and tgt in per_target_scale:
                s_t = float(per_target_scale[tgt])
                A = (1.0 + 0.15 * s_t) * A
                B = (1.0 - 0.15 * s_t) * B
            if target_weights and tgt in target_weights:
                tw = float(target_weights[tgt])
                s = float(max(0.0, tw)) ** 0.5
                A = s * A
                B = s * B
            # Optional norm cap on A/B (Frobenius) to avoid runaway scales
            try:
                if map_cap is not None and float(map_cap) > 0:
                    capv = float(map_cap)
                    nA = float(np.linalg.norm(A)) if A.size > 0 else 0.0
                    if nA > capv and nA > 0:
                        A = (capv / nA) * A
                    nB = float(np.linalg.norm(B)) if B.size > 0 else 0.0
                    if nB > capv and nB > 0:
                        B = (capv / nB) * B
            except Exception:
                pass
            e: Dict[str, np.ndarray] = {"A": A, "B": B, "gate": np.array([gate], dtype=np.float32)}
            if learn_bias:
                e["bias"] = np.zeros((int(d_out),), dtype=np.float32)
            layer_state[tgt] = e
        layers.append(layer_state)

    return {"layers": layers, "rank": rank, "d_model": d_model, "targets": targets, "gates": np.array(gates, dtype=np.float32)}

def generate_lora_from_embedding_torch(
    z: np.ndarray,
    *,
    d_model: int,
    num_layers: int,
    rank: int = 8,
    seed: int = 0,
    targets: Optional[List[str]] = None,
    target_shapes: Optional[Dict[str, Tuple[int, int]]] = None,
    einsum_opt: str = "auto",
    layer_gate: str = "zmean",
    target_weights: Optional[Dict[str, float]] = None,
    map_cap: Optional[float] = None,
) -> Dict[str, List[Dict[str, Dict[str, np.ndarray]]]]:
    import torch  # local import to avoid hard dep when unused

    if targets is None:
        targets = ["q_proj", "o_proj", "up_proj"]


            # Optional norm cap on A/B (Frobenius) to avoid runaway scales
            try:
                if map_cap is not None and float(map_cap) > 0:
                    capv = float(map_cap)
                    nA = float(np.linalg.norm(A)) if A.size > 0 else 0.0
                    if nA > capv and nA > 0:
                        A = (capv / nA) * A
                    nB = float(np.linalg.norm(B)) if B.size > 0 else 0.0
                    if nB > capv and nB > 0:
                        B = (capv / nB) * B
            except Exception:
                pass
            e: Dict[str, np.ndarray] = {"A": A, "B": B, "gate": np.array([gate], dtype=np.float32)}
            if learn_bias:
                e["bias"] = np.zeros((int(d_out),), dtype=np.float32)
            layer_state[tgt] = e
        layers.append(layer_state)

    return {"layers": layers, "rank": rank, "d_model": d_model, "targets": targets, "gates": np.array(gates, dtype=np.float32)}

def generate_lora_from_embedding_torch(
    z: np.ndarray,
    *,
    d_model: int,
    num_layers: int,
    rank: int = 8,
    seed: int = 0,
    targets: Optional[List[str]] = None,
    target_shapes: Optional[Dict[str, Tuple[int, int]]] = None,
    einsum_opt: str = "auto",
    layer_gate: str = "zmean",
    target_weights: Optional[Dict[str, float]] = None,
    map_cap: Optional[float] = None,
) -> Dict[str, List[Dict[str, Dict[str, np.ndarray]]]]:
    import torch  # local import to avoid hard dep when unused

    if targets is None:
        targets = ["q_proj", "o_proj", "up_proj"]

    zt = torch.from_numpy(z.astype(np.float32))
    try:
        nz = float(torch.linalg.vector_norm(zt).item())
        if nz > 0:
            zt = (zt / nz).to(torch.float32)
    except Exception:
        pass
    layers: List[Dict[str, Dict[str, np.ndarray]]] = []
    gates: List[float] = []
    # Global seeding for reproducibility in case external torch ops run
    try:
        torch.manual_seed(int(seed))
    except Exception:
        pass

    for layer_idx in range(num_layers):
        layer_state: Dict[str, Dict[str, np.ndarray]] = {}
        key = int((_stable_hash(f"layer:{layer_idx}", seed) ^ _stable_hash(str(float(z[0])), seed + 7)) & ((1 << 31) - 1))
        frac = float(layer_idx) / float(max(1, num_layers - 1))
        if layer_gate == "cosine":
            gate = float(0.5 * (1.0 - math.cos(math.pi * frac)))
        elif layer_gate == "hump":
            gate = float(max(0.0, math.sin(math.pi * frac)))
        elif layer_gate == "linear":
            gate = float(frac)
        else:
            gate = float((np.tanh(z[(layer_idx * 13) % len(z)]) + 1.0) * 0.5)
        gates.append(gate)
        # deterministic torch RNG
        gen = torch.Generator(device="cpu")
        gen.manual_seed(key)
        mlp_seed = key ^ _stable_hash("mlp_pair", seed)
        up_pair: Optional[Tuple[torch.Tensor, torch.Tensor]] = None
        # Precompute per-target scale (always-on)
        per_target_scale: Dict[str, float] = {}
        try:
            seg_len = max(8, int(len(zt) // max(1, len(targets))))
            for idx_t, tgt in enumerate(targets):
                start = (layer_idx * 17 + idx_t * seg_len) % len(zt)
                idxs = (torch.arange(seg_len) + start) % len(zt)
                seg = zt[idxs]
                per_target_scale[tgt] = float(torch.tanh(seg.mean() * 1.5).item())
        except Exception:
            per_target_scale = {}
        for tgt in targets:
            if target_shapes and tgt in target_shapes:
                d_out, d_in = target_shapes[tgt]
            else:
                d_out, d_in = d_model, d_model
            if tgt in ("up_proj", "gate_proj"):
                if up_pair is None:
                    A = (torch.rand((int(d_out), rank), generator=gen) * 2 - 1).to(torch.float32)
                    B = (torch.rand((rank, int(d_in)), generator=gen) * 2 - 1).to(torch.float32)
                    up_pair = (A, B)
                else:
                    A, B = up_pair
            elif tgt == "down_proj":
                # reuse generator but distinct suffix for down
                A = (torch.rand((int(d_out), rank), generator=gen) * 2 - 1).to(torch.float32)
                B = (torch.rand((rank, int(d_in)), generator=gen) * 2 - 1).to(torch.float32)
            else:
                A = (torch.rand((int(d_out), rank), generator=gen) * 2 - 1).to(torch.float32)
                B = (torch.rand((rank, int(d_in)), generator=gen) * 2 - 1).to(torch.float32)
            if rank > 0:
                A = A * (1.0 / float(max(1.0, math.sqrt(float(rank)))))
            # Modulate by contraction with a deterministic kernel vector derived from z
            start = (layer_idx * 31) % len(z)
            idx = (np.arange(32) + start) % len(z)
            seg = zt[idx]
            w = torch.sin(torch.linspace(0, 3.14159, steps=32))
            if _contract is not None and einsum_opt:
                alpha = torch.tanh(_contract("i,i->", seg, w, optimize=einsum_opt) / 8.0)
            else:
                alpha = torch.tanh((seg * w).sum() / 8.0)
            A = (1.0 + float(alpha.item()) * gate) * A
            B = (1.0 - float(alpha.item()) * gate) * B
            if per_target_scale and tgt in per_target_scale:
                s_t = float(per_target_scale[tgt])
                A = (1.0 + 0.15 * s_t) * A
                B = (1.0 - 0.15 * s_t) * B
            if target_weights and tgt in target_weights:
                tw = float(target_weights[tgt])
                s = float(max(0.0, tw)) ** 0.5
                A = (s * A)
                B = (s * B)
            # Optional norm cap (Frobenius) on A/B
            try:
                if map_cap is not None and float(map_cap) > 0:
                    capv = float(map_cap)
                    nA = float(torch.linalg.matrix_norm(A).item()) if A.numel() > 0 else 0.0
                    if nA > capv and nA > 0:
                        A = (capv / nA) * A
                    nB = float(torch.linalg.matrix_norm(B).item()) if B.numel() > 0 else 0.0
                    if nB > capv and nB > 0:
                        B = (capv / nB) * B
            except Exception:
                pass
            e: Dict[str, np.ndarray] = {
                "A": A.numpy().astype(np.float32),
                "B": B.numpy().astype(np.float32),
                "gate": np.array([gate], dtype=np.float32),
            }
            # Torch path does not currently support learn_bias flag; add zero bias for parity if requested via target_weights special key later if needed
            layer_state[tgt] = e
        layers.append(layer_state)

    return {"layers": layers, "rank": rank, "d_model": d_model, "targets": targets, "gates": np.array(gates, dtype=np.float32)}

def save_npz(out_dir: str, *, embedding: Dict[str, np.ndarray], adapters: Dict[str, Any], manifest: Dict[str, Any]) -> None:
    os.makedirs(out_dir, exist_ok=True)
    np.savez_compressed(os.path.join(out_dir, "embedding.npz"), **embedding)
    # flatten adapter arrays
    flat: Dict[str, np.ndarray] = {}
    for i, layer in enumerate(adapters["layers"]):
        for name, tensors in layer.items():
            flat[f"L{i}.{name}.A"] = tensors["A"]
            flat[f"L{i}.{name}.B"] = tensors["B"]
            gate_val = float(tensors["gate"][0]) if isinstance(tensors.get("gate"), np.ndarray) else float(tensors.get("gate", 0.0))
            flat[f"L{i}.{name}.gate"] = np.array(gate_val, dtype=np.float32)
            if "bias" in tensors and isinstance(tensors["bias"], np.ndarray):
                flat[f"L{i}.{name}.bias"] = tensors["bias"]
    np.savez_compressed(os.path.join(out_dir, "adapters.npz"), **flat)
    open(os.path.join(out_dir, "manifest.json"), "w", encoding="utf-8").write(json.dumps(manifest, indent=2))

def load_adapters_npz(path: str) -> Dict[str, List[Dict[str, Dict[str, np.ndarray]]]]:
    data = np.load(path)


                    nA = float(torch.linalg.matrix_norm(A).item()) if A.numel() > 0 else 0.0
                    if nA > capv and nA > 0:
                        A = (capv / nA) * A
                    nB = float(torch.linalg.matrix_norm(B).item()) if B.numel() > 0 else 0.0
                    if nB > capv and nB > 0:
                        B = (capv / nB) * B
            except Exception:
                pass
            e: Dict[str, np.ndarray] = {
                "A": A.numpy().astype(np.float32),
                "B": B.numpy().astype(np.float32),
                "gate": np.array([gate], dtype=np.float32),
            }
            # Torch path does not currently support learn_bias flag; add zero bias for parity if requested via target_weights special key later if needed
            layer_state[tgt] = e
        layers.append(layer_state)

    return {"layers": layers, "rank": rank, "d_model": d_model, "targets": targets, "gates": np.array(gates, dtype=np.float32)}

def save_npz(out_dir: str, *, embedding: Dict[str, np.ndarray], adapters: Dict[str, Any], manifest: Dict[str, Any]) -> None:
    os.makedirs(out_dir, exist_ok=True)
    np.savez_compressed(os.path.join(out_dir, "embedding.npz"), **embedding)
    # flatten adapter arrays
    flat: Dict[str, np.ndarray] = {}
    for i, layer in enumerate(adapters["layers"]):
        for name, tensors in layer.items():
            flat[f"L{i}.{name}.A"] = tensors["A"]
            flat[f"L{i}.{name}.B"] = tensors["B"]
            gate_val = float(tensors["gate"][0]) if isinstance(tensors.get("gate"), np.ndarray) else float(tensors.get("gate", 0.0))
            flat[f"L{i}.{name}.gate"] = np.array(gate_val, dtype=np.float32)
            if "bias" in tensors and isinstance(tensors["bias"], np.ndarray):
                flat[f"L{i}.{name}.bias"] = tensors["bias"]
    np.savez_compressed(os.path.join(out_dir, "adapters.npz"), **flat)
    open(os.path.join(out_dir, "manifest.json"), "w", encoding="utf-8").write(json.dumps(manifest, indent=2))

def load_adapters_npz(path: str) -> Dict[str, List[Dict[str, Dict[str, np.ndarray]]]]:
    data = np.load(path)
    # infer indices
    layers: Dict[int, Dict[str, Dict[str, np.ndarray]]] = {}
    for key in data.files:
        # L{idx}.{name}.{A|B|gate}
        parts = key.split(".")
        if len(parts) != 3:
            continue
        lid = int(parts[0][1:])
        name = parts[1]
        kind = parts[2]
        layers.setdefault(lid, {}).setdefault(name, {})[kind] = data[key]
    ordered = [layers[i] for i in sorted(layers.keys())]
    return {"layers": ordered}


    return {"layers": layers, "rank": rank, "d_model": d_model, "targets": targets, "gates": np.array(gates, dtype=np.float32)}

def save_npz(out_dir: str, *, embedding: Dict[str, np.ndarray], adapters: Dict[str, Any], manifest: Dict[str, Any]) -> None:
    os.makedirs(out_dir, exist_ok=True)
    np.savez_compressed(os.path.join(out_dir, "embedding.npz"), **embedding)
    # flatten adapter arrays
    flat: Dict[str, np.ndarray] = {}
    for i, layer in enumerate(adapters["layers"]):
        for name, tensors in layer.items():
            flat[f"L{i}.{name}.A"] = tensors["A"]
            flat[f"L{i}.{name}.B"] = tensors["B"]
            gate_val = float(tensors["gate"][0]) if isinstance(tensors.get("gate"), np.ndarray) else float(tensors.get("gate", 0.0))
            flat[f"L{i}.{name}.gate"] = np.array(gate_val, dtype=np.float32)
            if "bias" in tensors and isinstance(tensors["bias"], np.ndarray):
                flat[f"L{i}.{name}.bias"] = tensors["bias"]
    np.savez_compressed(os.path.join(out_dir, "adapters.npz"), **flat)
    open(os.path.join(out_dir, "manifest.json"), "w", encoding="utf-8").write(json.dumps(manifest, indent=2))

def load_adapters_npz(path: str) -> Dict[str, List[Dict[str, Dict[str, np.ndarray]]]]:
    data = np.load(path)
    # infer indices
    layers: Dict[int, Dict[str, Dict[str, np.ndarray]]] = {}
    for key in data.files:
        # L{idx}.{name}.{A|B|gate}
        parts = key.split(".")
        if len(parts) != 3:
            continue
        lid = int(parts[0][1:])
        name = parts[1]
        kind = parts[2]
        layers.setdefault(lid, {}).setdefault(name, {})[kind] = data[key]
    ordered = [layers[i] for i in sorted(layers.keys())]
    return {"layers": ordered}


modules.mixing

# delta-cap guard (lift from modular.py)
import os
import re
from typing import Any, Dict, List, Optional, Tuple

import numpy as np
import torch

from blocks.targets import targets_map
from blocks.utils import getattr_nested

def register_hook_mixed_adapters(
    model: Any,
    base_layers: List[Dict[str, Dict[str, np.ndarray]]],
    sub_layers: Optional[List[Dict[str, Dict[str, np.ndarray]]]],
    *,
    alpha_star: float,
    g_sub: float,
    rank: int,
    beta: float,
    target_weights: Optional[Dict[str, float]] = None,
    backend: str = "local",
    layer_multipliers: Optional[List[float]] = None,
    per_target_keep: Optional[Dict[str, int]] = None,
    per_target_keep_layers: Optional[List[Dict[str, int]]] = None,
    delta_cap: Optional[float] = None,
):
    tmap = targets_map(backend)
    target_weights = target_weights or {}
    per_target_keep = per_target_keep or {}
    per_target_keep_layers = per_target_keep_layers or []

    def _mat(A: np.ndarray, B: np.ndarray, *, keep: Optional[int] = None) -> torch.Tensor:
        if isinstance(keep, int) and keep > 0:
            k = int(min(keep, A.shape[1], B.shape[0]))
            if k > 0:
                At = torch.from_numpy(A[:, :k]).to(torch.float32)
                Bt = torch.from_numpy(B[:k, :]).to(torch.float32)
                return At @ Bt
        return torch.from_numpy(A).to(torch.float32) @ torch.from_numpy(B).to(torch.float32)

    L = max(len(base_layers or []), len(sub_layers or []))
    scale = float(alpha_star) / float(max(1, int(rank)))
    # Store CPU-offloaded records of applied deltas to avoid holding large GPU tensors
    # Each record: (param, cpu_delta, mode, extra)
    #   mode == "full"  -> subtract full delta
    #   mode == "slice" -> subtract into row slice [start:end]
    applied: List[Tuple[torch.nn.Parameter, torch.Tensor, str, Tuple[int, int]]] = []
    # layers list
    try:
        layers = list(getattr(getattr(model, "model", model), "layers"))
    except Exception:
        try:
            layers = list(getattr(model, "blocks"))
        except Exception:
            layers = []

    def _cap_delta_if_needed(w: torch.nn.Parameter, d: torch.Tensor) -> torch.Tensor:
        try:
            # Prefer explicit parameter; else env fallbacks
            if delta_cap is not None:
                cap = float(delta_cap)
            else:
                cap_env = (os.environ.get("PCA_DELTA_CAP") or os.environ.get("REPO_ADAPTER_DELTA_CAP") or "").strip()
                cap = float(cap_env) if cap_env else 0.0
        except Exception:
            cap = 0.0
        if cap is None or cap <= 0:
            return d
        try:
            wn = float(w.data.norm().item())
            dn = float(d.norm().item())
            if dn > 0 and wn > 0:
                limit = cap * wn
                if dn > limit:
                    s = float(limit / max(dn, 1e-12))
                    return (s * d)
        except Exception:
            return d
        return d

    for i in range(L):
        if i >= len(layers):
            break
        layer = layers[i]
        lm = 1.0
        try:
            if layer_multipliers is not None and i < len(layer_multipliers):
                lm = float(layer_multipliers[i])
        except Exception:
            lm = 1.0
        # Build deltas for this layer only on CPU, apply, then free
        base = base_layers[i] if i < len(base_layers or []) else None
        sub = (sub_layers[i] if (sub_layers is not None and i < len(sub_layers)) else None)
        deltas_cpu: Dict[str, torch.Tensor] = {}
        for name in tmap.keys():
            acc: Optional[torch.Tensor] = None
            # layer-specific keep overrides global keep
            keep_layer = None
            try:
                if 0 <= i < len(per_target_keep_layers):
                    keep_layer = int(per_target_keep_layers[i].get(name, 0))
            except Exception:
                keep_layer = None
            keep_global = int(per_target_keep.get(name, 0)) if per_target_keep else 0
            keep_eff = keep_layer if (keep_layer and keep_layer > 0) else (keep_global if keep_global > 0 else None)
            if base is not None and name in base:
                try:
                    acc = _mat(base[name]["A"], base[name]["B"], keep=keep_eff).to(torch.float32)
                except Exception:
                    acc = None
            if sub is not None and name in sub:
                try:
                    sub_m = _mat(sub[name]["A"], sub[name]["B"], keep=keep_eff).to(torch.float32)
                    acc = sub_m if acc is None else ((1.0 - float(g_sub)) * acc + float(g_sub) * sub_m)
                except Exception:
                    pass
            if acc is not None:
                tw = float(target_weights.get(name, 1.0))
                deltas_cpu[name] = (lm * tw * acc).contiguous()
        for short, rel in tmap.items():
            if short not in deltas_cpu:
                continue
            try:
                mod = getattr_nested(layer, rel)
                w: torch.nn.Parameter = getattr(mod, "weight")  # type: ignore[assignment]
                d_cpu = (scale * deltas_cpu[short])
                d = d_cpu.to(w.device, dtype=w.dtype)
                d = _cap_delta_if_needed(w, d)
                if rel.endswith("mlp.w_in") and short in ("gate_proj", "up_proj"):
                    half = int(w.shape[0] // 2)
                    if short == "gate_proj":
                        part = d[:half, :] if d.shape[0] == w.shape[0] else d
                        part = _cap_delta_if_needed(w, part)
                        w.data[:half, :].add_(part)
                        applied.append((w, d_cpu[:half, :].detach().to("cpu"), "slice", (0, half)))
                    else:
                        part = d[-half:, :] if d.shape[0] == w.shape[0] else d
                        part = _cap_delta_if_needed(w, part)
                        w.data[half:, :].add_(part)
                        applied.append((w, d_cpu[-half:, :].detach().to("cpu"), "slice", (half, w.shape[0])))
                else:
                    w.data.add_(d)
                    applied.append((w, d_cpu.detach().to("cpu"), "full", (0, 0)))
            except Exception:
                continue
        del deltas_cpu

    class _Applied:
        def __init__(self, records: List[Tuple[torch.nn.Parameter, torch.Tensor, str, Tuple[int, int]]]):
            self._records = records
            self._removed = False

        def remove(self) -> None:
            if self._removed:
                return
            for w, d_cpu, mode, extra in self._records:
                try:
                    if mode == "full":
                        d_dev = d_cpu.to(device=w.device, dtype=w.dtype)
                        w.data.sub_(d_dev)
                    elif mode == "slice":
                        start, end = extra
                        d_dev = d_cpu.to(device=w.device, dtype=w.dtype)
                        w.data[start:end, :].sub_(d_dev)
                except Exception:
                    pass
            self._removed = True

    return [_Applied(applied)]


    model: Any,
    base_layers: List[Dict[str, Dict[str, np.ndarray]]],
    sub_layers: Optional[List[Dict[str, Dict[str, np.ndarray]]]],
    *,
    alpha_star: float,
    g_sub: float,
    rank: int,
    beta: float,
    target_weights: Optional[Dict[str, float]] = None,
    backend: str = "local",
    layer_multipliers: Optional[List[float]] = None,
    per_target_keep: Optional[Dict[str, int]] = None,
    per_target_keep_layers: Optional[List[Dict[str, int]]] = None,
    delta_cap: Optional[float] = None,
):
    tmap = targets_map(backend)
    target_weights = target_weights or {}
    per_target_keep = per_target_keep or {}
    per_target_keep_layers = per_target_keep_layers or []

    def _mat(A: np.ndarray, B: np.ndarray, *, keep: Optional[int] = None) -> torch.Tensor:
        if isinstance(keep, int) and keep > 0:
            k = int(min(keep, A.shape[1], B.shape[0]))
            if k > 0:
                At = torch.from_numpy(A[:, :k]).to(torch.float32)
                Bt = torch.from_numpy(B[:k, :]).to(torch.float32)
                return At @ Bt
        return torch.from_numpy(A).to(torch.float32) @ torch.from_numpy(B).to(torch.float32)

    L = max(len(base_layers or []), len(sub_layers or []))
    scale = float(alpha_star) / float(max(1, int(rank)))
    # Store CPU-offloaded records of applied deltas to avoid holding large GPU tensors
    # Each record: (param, cpu_delta, mode, extra)
    #   mode == "full"  -> subtract full delta
    #   mode == "slice" -> subtract into row slice [start:end]
    applied: List[Tuple[torch.nn.Parameter, torch.Tensor, str, Tuple[int, int]]] = []
    # layers list
    try:
        layers = list(getattr(getattr(model, "model", model), "layers"))
    except Exception:
        try:
            layers = list(getattr(model, "blocks"))
        except Exception:
            layers = []

    def _cap_delta_if_needed(w: torch.nn.Parameter, d: torch.Tensor) -> torch.Tensor:
        try:
            # Prefer explicit parameter; else env fallbacks


                Bt = torch.from_numpy(B[:k, :]).to(torch.float32)
                return At @ Bt
        return torch.from_numpy(A).to(torch.float32) @ torch.from_numpy(B).to(torch.float32)

    L = max(len(base_layers or []), len(sub_layers or []))
    scale = float(alpha_star) / float(max(1, int(rank)))
    # Store CPU-offloaded records of applied deltas to avoid holding large GPU tensors
    # Each record: (param, cpu_delta, mode, extra)
    #   mode == "full"  -> subtract full delta
    #   mode == "slice" -> subtract into row slice [start:end]
    applied: List[Tuple[torch.nn.Parameter, torch.Tensor, str, Tuple[int, int]]] = []
    # layers list
    try:
        layers = list(getattr(getattr(model, "model", model), "layers"))
    except Exception:
        try:
            layers = list(getattr(model, "blocks"))
        except Exception:
            layers = []

    def _cap_delta_if_needed(w: torch.nn.Parameter, d: torch.Tensor) -> torch.Tensor:
        try:
            # Prefer explicit parameter; else env fallbacks
            if delta_cap is not None:
                cap = float(delta_cap)
            else:
                cap_env = (os.environ.get("PCA_DELTA_CAP") or os.environ.get("REPO_ADAPTER_DELTA_CAP") or "").strip()
                cap = float(cap_env) if cap_env else 0.0
        except Exception:
            cap = 0.0
        if cap is None or cap <= 0:
            return d
        try:
            wn = float(w.data.norm().item())
            dn = float(d.norm().item())
            if dn > 0 and wn > 0:
                limit = cap * wn
                if dn > limit:
                    s = float(limit / max(dn, 1e-12))
                    return (s * d)
        except Exception:
            return d
        return d

    for i in range(L):
        if i >= len(layers):
            break
        layer = layers[i]
        lm = 1.0
        try:
            if layer_multipliers is not None and i < len(layer_multipliers):
                lm = float(layer_multipliers[i])
        except Exception:
            lm = 1.0
        # Build deltas for this layer only on CPU, apply, then free
        base = base_layers[i] if i < len(base_layers or []) else None
        sub = (sub_layers[i] if (sub_layers is not None and i < len(sub_layers)) else None)
        deltas_cpu: Dict[str, torch.Tensor] = {}
        for name in tmap.keys():
            acc: Optional[torch.Tensor] = None
            # layer-specific keep overrides global keep
            keep_layer = None
            try:


                d = _cap_delta_if_needed(w, d)
                if rel.endswith("mlp.w_in") and short in ("gate_proj", "up_proj"):
                    half = int(w.shape[0] // 2)
                    if short == "gate_proj":
                        part = d[:half, :] if d.shape[0] == w.shape[0] else d
                        part = _cap_delta_if_needed(w, part)
                        w.data[:half, :].add_(part)
                        applied.append((w, d_cpu[:half, :].detach().to("cpu"), "slice", (0, half)))
                    else:
                        part = d[-half:, :] if d.shape[0] == w.shape[0] else d
                        part = _cap_delta_if_needed(w, part)
                        w.data[half:, :].add_(part)
                        applied.append((w, d_cpu[-half:, :].detach().to("cpu"), "slice", (half, w.shape[0])))
                else:
                    w.data.add_(d)
                    applied.append((w, d_cpu.detach().to("cpu"), "full", (0, 0)))
            except Exception:
                continue
        del deltas_cpu

    class _Applied:
        def __init__(self, records: List[Tuple[torch.nn.Parameter, torch.Tensor, str, Tuple[int, int]]]):
            self._records = records
            self._removed = False

        def remove(self) -> None:
            if self._removed:
                return
            for w, d_cpu, mode, extra in self._records:
                try:
                    if mode == "full":
                        d_dev = d_cpu.to(device=w.device, dtype=w.dtype)
                        w.data.sub_(d_dev)
                    elif mode == "slice":
                        start, end = extra
                        d_dev = d_cpu.to(device=w.device, dtype=w.dtype)
                        w.data[start:end, :].sub_(d_dev)
                except Exception:
                    pass
            self._removed = True

    return [_Applied(applied)]


                if rel.endswith("mlp.w_in") and short in ("gate_proj", "up_proj"):
                    half = int(w.shape[0] // 2)
                    if short == "gate_proj":
                        part = d[:half, :] if d.shape[0] == w.shape[0] else d
                        part = _cap_delta_if_needed(w, part)
                        w.data[:half, :].add_(part)
                        applied.append((w, d_cpu[:half, :].detach().to("cpu"), "slice", (0, half)))
                    else:
                        part = d[-half:, :] if d.shape[0] == w.shape[0] else d
                        part = _cap_delta_if_needed(w, part)
                        w.data[half:, :].add_(part)
                        applied.append((w, d_cpu[-half:, :].detach().to("cpu"), "slice", (half, w.shape[0])))
                else:
                    w.data.add_(d)
                    applied.append((w, d_cpu.detach().to("cpu"), "full", (0, 0)))
            except Exception:
                continue
        del deltas_cpu

    class _Applied:
        def __init__(self, records: List[Tuple[torch.nn.Parameter, torch.Tensor, str, Tuple[int, int]]]):
            self._records = records
            self._removed = False

        def remove(self) -> None:
            if self._removed:
                return
            for w, d_cpu, mode, extra in self._records:
                try:
                    if mode == "full":
                        d_dev = d_cpu.to(device=w.device, dtype=w.dtype)
                        w.data.sub_(d_dev)
                    elif mode == "slice":
                        start, end = extra
                        d_dev = d_cpu.to(device=w.device, dtype=w.dtype)
                        w.data[start:end, :].sub_(d_dev)
                except Exception:
                    pass
            self._removed = True

    return [_Applied(applied)]