js/ir.js

// ir.js — Architecture IR model + a tiny client-side resolver.
//
// The IR ("architecture-template-v0") is a config-parametric, semantic
// description of a Transformers architecture. This module turns the raw
// artifact JSON into a unified node registry + containment tree that the
// layout and rendering layers can consume, and resolves symbolic repeat
// counts against a config.json-like object.
//
// We deliberately consume the IR as-is and do not reinterpret it: the tree
// comes from `children` links + repeat `body`, edges are kept verbatim.

const KIND_LABELS = {
  model: "Model",
  embedding: "Embedding",
  encoder: "Encoder",
  decoder: "Decoder",
  attention: "Attention",
  cross_attention: "Cross-Attention",
  feed_forward: "Feed-Forward",
  projection: "Projection",
  moe: "Mixture-of-Experts",
  lm_head: "LM Head",
  normalization: "Normalization",
  position: "Position",
  pooler: "Pooler",
  transformer_block: "Transformer Block",
  repeated_container: "Container",
  symbolic_repeat: "Repeat",
};

// Split "LlamaDecoderLayer" -> "Llama Decoder Layer"; leave acronyms grouped.
function humanizeClassName(name) {
  if (!name) return "";
  return name
    .replace(/([a-z0-9])([A-Z])/g, "$1 $2")
    .replace(/([A-Z]+)([A-Z][a-z])/g, "$1 $2")
    .trim();
}

// A readable label for a node id like "decoder_layer.self_attn".
function humanizeId(id) {
  const leaf = String(id).split(".").pop();
  return leaf
    .replace(/[_]+/g, " ")
    .replace(/\b\w/g, (c) => c.toUpperCase());
}

// Evaluate a repeat count expression such as "config.num_hidden_layers"
// against a config-fields object. Falls back to the baked-in count.
// Supports simple arithmetic over config.<field> tokens and literals.
function evalCountExpr(expr, fields, fallback) {
  if (expr == null) return fallback;
  if (typeof expr === "number") return expr;
  const raw = String(expr).trim();
  // Substitute config.<ident> with its numeric value.
  let substituted = raw.replace(/config\.([A-Za-z_][A-Za-z0-9_]*)/g, (_, key) => {
    const v = fields ? fields[key] : undefined;
    return typeof v === "number" ? String(v) : "NaN";
  });
  // Only allow a safe arithmetic subset.
  if (/^[\d\s+\-*/().]+$/.test(substituted)) {
    try {
      // eslint-disable-next-line no-new-func
      const val = Function(`"use strict";return (${substituted});`)();
      if (typeof val === "number" && Number.isFinite(val)) {
        return Math.max(0, Math.round(val));
      }
    } catch (_) {
      /* fall through */
    }
  }
  return fallback;
}

// The config block exposes `referenced_fields` (authoritative parametric
// surface — what repeat counts / shapes reference) and `salient_fields` (a
// curated scalar whitelist for display). Merge into one name→value map, with
// referenced_fields winning on conflict.
function collectConfigFields(cfg) {
  const out = {};
  if (!cfg) return out;
  for (const key of ["salient_fields", "referenced_fields"]) {
    const v = cfg[key];
    if (v && typeof v === "object" && !Array.isArray(v)) Object.assign(out, v);
  }
  return out;
}

class IRModel {
  constructor(raw) {
    this.raw = raw;
    this.modelType = raw.model_type || (raw.config && raw.config.model_type) || "model";
    this.schemaVersion = raw.schema_version;
    // Slimmed provenance: model_class/model_module/config_class/config_module.
    this.provenance = raw.provenance || {};
    // Parametric surface: referenced_fields + salient_fields (no full config).
    this.configFields = collectConfigFields(raw.config);
    this.configClass = raw.config && raw.config.class_name;
    // Semantic facts, observed dataflow, modular inheritance.
    this.architecture = raw.architecture || null;
    this.dataflow = raw.dataflow || null;
    // Modular: `extends` is the parent (null = standalone); `patches` the diffs.
    this.extends = raw.extends || null;
    this.patches = Array.isArray(raw.patches) ? raw.patches : [];
    this.capabilities = raw.capabilities || null; // backends / schedule / task heads / TP

    this.nodes = new Map();      // id -> node
    this.repeatsById = new Map();
    this.edges = [];
    this.pseudoIds = new Set();  // "input:*" / "state:*" pseudo endpoints

    this._index();
    this._buildTree();
    this._indexEdges();
    this._synthesizeLmHead();
  }

  // The base-model IR stops at the final hidden state — it has no LM head.
  // But whether the output projection is a *separate* set of weights or shares
  // the input embedding's is load-bearing (param count, checkpoint layout), so
  // we synthesize an `lm_head` node from the word embedding (hidden → vocab).
  // It is always present in the tree; the layout shows it only when *untied*
  // and shows a "tied" badge on the embedding otherwise (see tiedEmbeddings()).
  _synthesizeLmHead() {
    if (!this.rootId) return;
    const root = this.nodes.get(this.rootId);
    if (!root) return;
    const wordEmb = [...this.nodes.values()].find((n) => this.isWordEmbedding(n));
    if (!wordEmb) return; // no vocab embedding → no LM-head concept (e.g. vision)

    // The flow sink at the top level: a direct child of root that receives data
    // but emits none (the final norm / decoder), i.e. what feeds the head.
    const kidSet = new Set(root.children);
    const topChild = (x) => {
      let cur = x;
      for (let g = 0; g < 64 && cur; g++) {
        if (kidSet.has(cur)) return cur;
        cur = (this.nodes.get(cur) || {}).parent;
      }
      return null;
    };
    const hasIn = new Set();
    const hasOut = new Set();
    for (const e of this.edges) {
      if (e.kind !== "data") continue;
      const s = topChild(e.source);
      const t = topChild(e.target);
      if (s && t && s !== t) {
        hasOut.add(s);
        hasIn.add(t);
      }
    }
    const sinks = root.children.filter((k) => hasIn.has(k) && !hasOut.has(k) && k !== wordEmb.id);
    const sinkId = sinks.length ? sinks[sinks.length - 1] : null;
    if (!sinkId) return;

    const a = wordEmb.attributes || {};
    const lm = {
      id: "lm_head",
      kind: "lm_head",
      className: "Linear",
      pathPattern: "lm_head",
      rawChildren: [],
      // hidden → vocab, mirrored from the (transposed) word embedding.
      attributes: {
        in_features: a.embedding_dim || "config.hidden_size",
        out_features: a.num_embeddings || "config.vocab_size",
      },
      nodeType: "synthetic",
      children: [],
      parent: this.rootId,
    };
    this.nodes.set(lm.id, lm);
    root.children.push(lm.id);
    this.lmHeadId = lm.id;
    this.lmHeadFrom = sinkId;
    this.edges.push({ kind: "data", source: sinkId, target: lm.id, _idx: this.edges.length, _synthetic: true });
  }

  _addRaw(list, nodeType) {
    (list || []).forEach((c) => {
      this.nodes.set(c.id, {
        id: c.id,
        kind: c.kind,
        className: c.class_name,
        pathPattern: c.path_pattern,
        rawChildren: c.children ?? [], // omitted on leaf nodes
        attributes: c.attributes || null, // per-component semantic facts (spec)
        nodeType, // 'component' | 'template'
        children: [],
        parent: null,
      });
    });
  }

  _index() {
    this._addRaw(this.raw.components, "component");
    this._addRaw(this.raw.templates, "template");

    (this.raw.repeats || []).forEach((r) => {
      const node = {
        id: r.id,
        kind: "symbolic_repeat",
        className: r.repeated_class_name,
        pathPattern: r.container_path_pattern,
        rawChildren: [r.body],
        nodeType: "repeat",
        children: [],
        parent: null,
        repeat: r,
      };
      this.nodes.set(r.id, node);
      this.repeatsById.set(r.id, node);
    });
  }

  _buildTree() {
    // Link children (ids may be components, templates or repeat ids).
    for (const node of this.nodes.values()) {
      node.rawChildren.forEach((cid) => {
        const child = this.nodes.get(cid);
        if (child) {
          node.children.push(cid);
          child.parent = node.id;
        }
      });
    }
    // A repeat's body template is drawn AS the repeat's group box, so chain
    // the body (and thus its whole subtree) up through the repeat node.
    this.bodyToRepeat = new Map();
    for (const r of this.repeatsById.values()) {
      const body = this.nodes.get(r.repeat.body);
      if (body) {
        body.parent = r.id;
        this.bodyToRepeat.set(body.id, r.id);
      }
    }

    // Root = a "model" node with no parent (fallback: first parentless node).
    let root =
      [...this.nodes.values()].find((n) => n.kind === "model" && !n.parent) ||
      [...this.nodes.values()].find((n) => !n.parent);
    this.rootId = root ? root.id : null;

    this._attachOrphans();
  }

  // Some artifacts (e.g. multimodal models) leave `model.children` empty and
  // don't link submodules via `children` — the hierarchy lives only in the
  // dotted `path_pattern`. Reattach any orphaned component/repeat by its path,
  // synthesizing intermediate container nodes as needed. Well-linked models are
  // untouched (they have no orphans).
  _attachOrphans() {
    if (!this.rootId) return;
    const rootPath = (this.nodes.get(this.rootId) || {}).pathPattern || "model";
    const pathToId = new Map();
    for (const [id, n] of this.nodes) {
      if (n.nodeType === "template") continue;
      if (n.pathPattern && !pathToId.has(n.pathPattern)) pathToId.set(n.pathPattern, id);
    }

    const ensureByPath = (path) => {
      if (pathToId.has(path)) return pathToId.get(path);
      if (!path || path === rootPath || !path.includes(".")) return this.rootId;
      const synth = {
        id: path,
        kind: "module",
        className: null,
        module: null,
        pathPattern: path,
        rawChildren: [],
        attributes: null,
        nodeType: "synthetic",
        children: [],
        parent: null,
      };
      this.nodes.set(path, synth);
      pathToId.set(path, path);
      const pid = ensureByPath(path.slice(0, path.lastIndexOf(".")));
      synth.parent = pid;
      const pn = this.nodes.get(pid);
      if (pn && !pn.children.includes(path)) pn.children.push(path);
      return path;
    };

    const link = (id, parentId) => {
      this.nodes.get(id).parent = parentId;
      const pn = this.nodes.get(parentId);
      if (pn && !pn.children.includes(id)) pn.children.push(id);
    };

    for (const [id, n] of [...this.nodes]) {
      if (id === this.rootId || n.parent || n.nodeType === "template" || n.nodeType === "synthetic") continue;
      let parentPath;
      if (n.nodeType === "repeat") {
        const cpp = n.repeat.container_path_pattern;
        parentPath = cpp && cpp.includes(".") ? cpp.slice(0, cpp.lastIndexOf(".")) : rootPath;
      } else {
        const p = n.pathPattern;
        parentPath = p && p.includes(".") ? p.slice(0, p.lastIndexOf(".")) : rootPath;
      }
      link(id, ensureByPath(parentPath));
    }
  }

  // A pseudo endpoint (input:/state:/…) — a graph source/sink that is not a
  // real module component.
  isPseudo(id) {
    return typeof id === "string" && id.includes(":") && !this.nodes.has(id);
  }

  _indexEdges() {
    const present = new Set();
    (this.raw.edges || []).forEach((e, idx) => {
      const edge = { ...e, _idx: idx };
      this.edges.push(edge);
      if (e.kind) present.add(e.kind);
      [e.source, e.target].forEach((ep) => {
        if (this.isPseudo(ep)) this.pseudoIds.add(ep);
      });
    });
    // Edge kinds actually present in this artifact, known ones first (stable
    // order), then any new/unknown kinds (e.g. cache, route, dataflow) sorted.
    const extra = [...present].filter((k) => !EDGE_KINDS.includes(k)).sort();
    this.edgeKinds = [...EDGE_KINDS.filter((k) => present.has(k)), ...extra];
    if (!this.edgeKinds.length) this.edgeKinds = [...EDGE_KINDS];
  }

  // --- Public helpers -------------------------------------------------------

  node(id) {
    return this.nodes.get(id);
  }

  kindLabel(kind) {
    return KIND_LABELS[kind] || humanizeId(kind || "");
  }

  // Structural children shown when a container is expanded.
  // A repeat's content is its body's children (the body itself is drawn as
  // the group box), so repeat + body collapse into one visual container.
  layoutChildrenOf(id) {
    const n = this.nodes.get(id);
    if (!n) return [];
    if (n.nodeType === "repeat") {
      const body = this.nodes.get(n.repeat.body);
      return body ? body.children.slice() : [];
    }
    return n.children.slice();
  }

  isContainer(id) {
    return this.layoutChildrenOf(id).length > 0;
  }

  // Resolve a repeat's count against a config-fields object.
  resolveCount(repeatNode, fields) {
    const r = repeatNode.repeat;
    return evalCountExpr(r.count_expr, fields, r.count);
  }

  // Label for a repeat block, e.g. "Decoder Layer × 32" (resolved against config).
  repeatLabel(repeatNode, fields) {
    const base = humanizeClassName(repeatNode.className) || humanizeId(repeatNode.id);
    return `${base} × ${this.resolveCount(repeatNode, fields)}`;
  }

  // A short display label for any node.
  label(id, fields) {
    const n = this.nodes.get(id);
    if (!n) {
      // Pseudo endpoint like "input:attention_mask" / "state:kv_cache".
      const rest = typeof id === "string" && id.includes(":") ? id.slice(id.indexOf(":") + 1) : id;
      return humanizeId(rest);
    }
    if (n.nodeType === "repeat") return this.repeatLabel(n, fields);
    return humanizeClassName(n.className) || humanizeId(n.id);
  }

  // Observed-forward shapes for a component id: { in, out } or null. (Ordering
  // is not here — it lives in the data edges, which layout already uses.)
  nodeShapes(id) {
    return (this.dataflow && this.dataflow.shapes && this.dataflow.shapes[id]) || null;
  }

  // Whether input word embeddings are tied to the output (LM head) — shared
  // weights. Read from the (editable) config first, then the architecture block,
  // so flipping tie_word_embeddings in the config reflects live.
  tiedEmbeddings(fields) {
    const f = fields || this.configFields;
    if (f && Object.prototype.hasOwnProperty.call(f, "tie_word_embeddings")) return !!f.tie_word_embeddings;
    return !!(this.architecture && this.architecture.tie_word_embeddings);
  }

  // Is this node the token/word embedding (kind embedding sized to vocab)?
  // Excludes position / token-type / vision embeddings.
  isWordEmbedding(node) {
    return !!(
      node &&
      node.kind === "embedding" &&
      node.attributes &&
      node.attributes.num_embeddings === "config.vocab_size" &&
      node.children.length === 0
    );
  }

  // A Hub kernel that can augment this node: { name, repos } from the node's
  // `attributes.kernel` and the model's `capabilities.kernels` map, or null.
  nodeKernel(node) {
    const name = node && node.attributes && node.attributes.kernel;
    if (!name) return null;
    const map = this.capabilities && this.capabilities.kernels;
    return { name, repos: (map && map[name]) || [] };
  }

  // The concrete model class(es): a checkpoint's `config.architectures` when
  // present (authoritative — each checkpoint pins its own), else the base
  // model class from provenance. Not fabricated from task heads.
  architectureClasses(fields) {
    const arch = fields && fields.architectures;
    if (Array.isArray(arch) && arch.length) return arch;
    const mc = this.provenance && this.provenance.model_class;
    return mc ? [mc] : [];
  }

  // Geometry for a repeat block's stacked "deck" (shared by the renderer and the
  // layout so they agree). Consecutive same-pattern layers merge into one solid
  // band whose thickness ∝ run length (min so a lone `full` band still shows),
  // scaled to a depth cap. Returns { cards:[{pattern|null, off}], dw, dh } or null.
  deckGeometry(node, fields) {
    if (!node || node.nodeType !== "repeat") return null;
    const sched = this.scheduleForRepeat(node, fields);
    let runs;
    if (sched && sched.length) {
      runs = [];
      for (const p of sched) {
        const last = runs[runs.length - 1];
        if (last && last.p === p) last.len++;
        else runs.push({ p, len: 1 });
      }
    } else {
      const n = this.resolveCount(node, fields);
      if (!n || n < 2) return null;
      runs = [{ p: null, len: Math.min(n, 64) }]; // uniform → one neutral band, depth ∝ layers
    }
    const UNIT = 1.6;
    const MIN = 3; // a lone layer (e.g. one `full`) still gets a visible band
    const CAP = 48;
    const bands = runs.map((r) => ({ p: r.p, t: r.p ? Math.max(r.len * UNIT, MIN) : r.len * UNIT }));
    const total = bands.reduce((s, b) => s + b.t, 0);
    const scale = total > CAP ? CAP / total : 1;
    let off = 0;
    const cards = [];
    for (const b of bands) {
      off += b.t * scale;
      cards.push({ pattern: b.p, off });
    }
    return { cards, dw: Math.round(off * 0.6), dh: Math.round(off) };
  }

  // Tensor-parallel sharding plan for a container: which descendant projections
  // are column- vs row-parallel. Colwise splits output features (no collective);
  // rowwise splits input features (needs an all-reduce). Returns null if none.
  tpPlan(id) {
    const start = this.nodes.get(id);
    if (!start) return null;
    const col = [];
    const row = [];
    const visit = (nid) => {
      const n = this.nodes.get(nid);
      if (!n) return;
      const tp = n.attributes && n.attributes.tp;
      if (n.kind === "projection" && tp) (tp === "rowwise" ? row : col).push(nid.split(".").pop());
      n.children.forEach(visit);
    };
    start.children.forEach(visit);
    return col.length || row.length ? { colwise: col, rowwise: row } : null;
  }

  // Per-model modular diff magnitude, derived from patches (the graph-wide
  // number lives in modular_graph.json): member-list lengths + 3×new classes.
  diffSize() {
    let n = 0;
    for (const p of this.patches) {
      for (const bucket of [p.added, p.overridden, p.deleted]) {
        if (!bucket) continue;
        n += (bucket.methods ?? []).length + (bucket.attrs ?? []).length;
      }
      if (p.relation === "new") n += 3;
    }
    return n;
  }

  // The per-layer attention schedule, bound to a repeat block iff its length
  // matches that repeat's resolved layer count (so e.g. gemma3's 26-long
  // schedule attaches to the text decoder, not the 12-layer vision encoder).
  scheduleForRepeat(repeatNode, fields) {
    const sched = this.capabilities && this.capabilities.attention_schedule;
    if (!Array.isArray(sched) || !repeatNode || repeatNode.nodeType !== "repeat") return null;
    return sched.length === this.resolveCount(repeatNode, fields) ? sched : null;
  }

  // Compact summary of a schedule: counts + smallest repeating period, e.g.
  // "22 sliding · 4 full — 5×sliding + 1×full".
  scheduleSummary(schedule) {
    if (!Array.isArray(schedule) || !schedule.length) return "";
    const counts = new Map();
    for (const p of schedule) counts.set(p, (counts.get(p) || 0) + 1);
    const short = (p) => String(p).replace(/_attention$/, "");
    const totals = [...counts.entries()].map(([p, c]) => `${c} ${short(p)}`).join(" · ");
    // Smallest period whose repetition reproduces the whole schedule.
    let period = schedule.length;
    for (let k = 1; k <= schedule.length / 2; k++) {
      if (schedule.length % k) continue;
      if (schedule.every((v, i) => v === schedule[i % k])) { period = k; break; }
    }
    if (period < schedule.length) {
      const runs = [];
      for (let i = 0; i < period; ) {
        let j = i;
        while (j < period && schedule[j] === schedule[i]) j++;
        runs.push(`${j - i}×${short(schedule[i])}`);
        i = j;
      }
      return `${totals} — ${runs.join(" + ")}`;
    }
    return totals;
  }

  // Resolve a symbolic shape like ["B","S","config.hidden_size"] against the
  // config into "[B, S, 4096]". Non-config dims (B, S, ints) pass through.
  resolveShape(shape, fields) {
    if (!Array.isArray(shape)) return null;
    const dims = shape.map((d) => {
      const s = String(d);
      if (s.startsWith("config.")) {
        const key = s.slice("config.".length);
        const v = fields ? fields[key] : undefined;
        return v != null ? String(v) : key;
      }
      return s;
    });
    return `[${dims.join(", ")}]`;
  }

  // A compact one-line caption for a node built from its observed output shape
  // and semantic attributes (both from the IR) — e.g. "[B, S, 4096]" or
  // "head dim 128 · n heads 32". Returns null when there's nothing to show.
  nodeInfo(id, fields) {
    const parts = [];
    const shapes = this.nodeShapes(id);
    if (shapes && shapes.out) {
      const sh = this.resolveShape(shapes.out, fields);
      if (sh) parts.push(sh);
    }
    const n = this.nodes.get(id);
    if (n && n.attributes && typeof n.attributes === "object") {
      const attrs = n.attributes;
      const val = (v) => {
        if (typeof v === "string" && v.startsWith("config.")) {
          const rv = fields ? fields[v.slice("config.".length)] : undefined;
          return rv != null ? rv : v.slice("config.".length);
        }
        return v;
      };
      // Linear-projection dims render as "in → out".
      if (attrs.in_features !== undefined && attrs.out_features !== undefined) {
        parts.push(`${val(attrs.in_features)} → ${val(attrs.out_features)}`);
      }
      const a = [];
      for (const [k, v] of Object.entries(attrs)) {
        // tp is shown as the col/row glyph; kernel as the ⚡ badge — not text.
        if (k === "in_features" || k === "out_features" || k === "tp" || k === "kernel") continue;
        if (v === true) a.push(k.replace(/_/g, " "));
        else if (v === false || v == null) continue;
        else a.push(`${k.replace(/_/g, " ")} ${val(v)}`);
      }
      if (a.length) parts.push(a.slice(0, 3).join(" · "));
    }
    let s = parts.join("   ·   ");
    if (s.length > 42) s = s.slice(0, 41) + "…";
    return s || null;
  }

  // Edges incident to a node id (verbatim IR edges).
  edgesFor(id) {
    const incoming = this.edges.filter((e) => e.target === id);
    const outgoing = this.edges.filter((e) => e.source === id);
    return { incoming, outgoing };
  }
}

const EDGE_KINDS = [
  "data",
  "residual",
  "mask",
  "position",
  "cross_attention",
  "cache_read",
  "cache_write",
];

// Colour per attention pattern — shared by the SVG strip (graph.js) and the
// inspector's DOM cells (app.js) so they always match.
function patternColor(pattern) {
  const p = String(pattern || "");
  if (p.startsWith("sliding")) return "#f59e0b"; // sliding-window → amber
  if (p.startsWith("full") || p === "causal") return "#3b82f6"; // full/causal → blue
  if (p === "mamba") return "#a78bfa"; // SSM → purple
  if (p.startsWith("bidirectional")) return "#2dd4bf"; // encoder → teal
  return "#94a3b8"; // unknown → grey
}

export { IRModel, EDGE_KINDS, humanizeClassName, humanizeId, evalCountExpr, patternColor };