// 04 — the complete ArchitectureTemplate (all tiers together), lightly annotated.
// Faithful trimmed copy of artifacts/mistral.json (a modular, GQA, sliding-window decoder LLM).
// JSONC = JSON + // comments (not machine-parseable as-is); "…" marks omitted repetition.
// Regenerate the exact file: python utils/architecture_ir/generate_architecture_ir.py --architectures mistral --output-dir out
{
  "schema_version": "architecture-template-v0",
  "model_type": "mistral",

  // modular inheritance (diff_size metric lives in modular_graph.json)
  "extends": "llama",
  "patches": [
    { "relation": "inherits", "target_class": "MistralAttention", "component_kind": "attention",
      "parent_class": "LlamaAttention", "overridden": { "methods": ["__init__", "forward"] } },
    { "relation": "new", "target_class": "MistralForQuestionAnswering", "component_kind": null,
      "parent_class": "MistralPreTrainedModel" }
    // … other overridden classes (MistralMLP, MistralModel, …)
  ],

  // model-level facts (multimodal reads the text backbone)
  "architecture": {
    "view": "decoder", "family": "causal_lm", "attention_variant": "GQA",
    "positional": "rope", "is_moe": false, "mixer": "attention", "sliding_window": 4096, "tie_word_embeddings": false
  },

  // what it can do / run with
  "capabilities": {
    "attention_backends": ["eager", "sdpa", "flash_attention", "flex_attention"],
    "attention_patterns": ["sliding"], "attention_schedule": null,
    "task_heads": ["causal_lm", "question_answering", "sequence_classification", "token_classification"],
    "tensor_parallel": true,
    "kernels": { "RMSNorm": ["kernels-community/liger-kernels", "kernels-community/rmsnorm", "kernels-community/mlx_rmsnorm"] }
  },

  // config identity + the parametric surface (full config NOT serialized)
  "config": {
    "class_name": "MistralConfig", "module": "transformers.models.mistral.configuration_mistral", "model_type": "mistral",
    "referenced_fields": { "num_hidden_layers": 32 },
    "salient_fields": { "hidden_size": 4096, "intermediate_size": 14336, "vocab_size": 32000,
                        "num_attention_heads": 32, "num_key_value_heads": 8, "head_dim": 128,
                        "max_position_embeddings": 131072, "sliding_window": 4096, "hidden_act": "silu",
                        "tie_word_embeddings": false }
  },

  // nodes outside any repeat body (leaf nodes omit `children`)
  "components": [
    { "id": "model", "kind": "model", "class_name": "MistralModel", "path_pattern": "model",
      "children": ["embed_tokens", "decoder_layers", "norm", "rotary_emb"] },
    { "id": "embed_tokens", "kind": "embedding", "class_name": "Embedding", "path_pattern": "model.embed_tokens",
      "attributes": { "num_embeddings": "config.vocab_size", "embedding_dim": "config.hidden_size" } },
    { "id": "norm", "kind": "normalization", "class_name": "MistralRMSNorm", "path_pattern": "model.norm",
      "attributes": { "norm_type": "rms", "kernel": "RMSNorm" } },
    { "id": "rotary_emb", "kind": "position", "class_name": "MistralRotaryEmbedding", "path_pattern": "model.rotary_emb",
      "attributes": { "scheme": "rope", "rope_theta": 10000.0, "head_dim": "config.head_dim" } }
  ],

  // the repeated block body, serialized once ({i} = layer index)
  "templates": [
    { "id": "decoder_layer", "kind": "transformer_block", "class_name": "MistralDecoderLayer",
      "path_pattern": "model.layers.{i}",
      "children": ["decoder_layer.input_layernorm", "decoder_layer.self_attn",
                   "decoder_layer.post_attention_layernorm", "decoder_layer.mlp"] },
    { "id": "decoder_layer.self_attn", "kind": "attention", "class_name": "MistralAttention",
      "path_pattern": "model.layers.{i}.self_attn",
      "children": ["decoder_layer.self_attn.q_proj", "decoder_layer.self_attn.k_proj",
                   "decoder_layer.self_attn.v_proj", "decoder_layer.self_attn.o_proj"],
      "attributes": { "variant": "GQA", "n_heads": 32, "n_kv_heads": 8, "head_dim": 128,
                      "rope": true, "sliding_window": 4096, "pattern": "sliding" } },
    // GQA: q/o are hidden-sized; k/v are num_kv_heads*head_dim = 1024
    { "id": "decoder_layer.self_attn.q_proj", "kind": "projection", "class_name": "Linear",
      "path_pattern": "model.layers.{i}.self_attn.q_proj",
      "attributes": { "in_features": "config.hidden_size", "out_features": 4096, "tp": "colwise" } },
    { "id": "decoder_layer.self_attn.k_proj", "kind": "projection", "class_name": "Linear",
      "path_pattern": "model.layers.{i}.self_attn.k_proj",
      "attributes": { "in_features": "config.hidden_size", "out_features": 1024, "tp": "colwise" } },
    // … v_proj (1024, colwise), o_proj (4096 → config.hidden_size, rowwise)
    { "id": "decoder_layer.mlp", "kind": "feed_forward", "class_name": "MistralMLP",
      "path_pattern": "model.layers.{i}.mlp",
      "children": ["decoder_layer.mlp.gate_proj", "decoder_layer.mlp.up_proj", "decoder_layer.mlp.down_proj"],
      "attributes": { "hidden_size": 4096, "intermediate_size": 14336, "activation": "silu" } },
    { "id": "decoder_layer.mlp.gate_proj", "kind": "projection", "class_name": "Linear",
      "path_pattern": "model.layers.{i}.mlp.gate_proj",
      "attributes": { "in_features": "config.hidden_size", "out_features": "config.intermediate_size", "tp": "colwise" } },
    // … up_proj (colwise), down_proj (config.intermediate_size → config.hidden_size, rowwise)
    { "id": "decoder_layer.input_layernorm", "kind": "normalization", "class_name": "MistralRMSNorm",
      "path_pattern": "model.layers.{i}.input_layernorm", "attributes": { "norm_type": "rms", "kernel": "RMSNorm" } }
    // … post_attention_layernorm
  ],

  // 32 identical layers collapsed to one symbolic entry, count kept parametric
  "repeats": [
    { "id": "decoder_layers", "kind": "symbolic_repeat", "body": "decoder_layer",
      "count_expr": "config.num_hidden_layers", "count": 32, "count_source": "config", "index_symbol": "i",
      "container_path_pattern": "model.layers", "item_path_pattern": "model.layers.{i}",
      "repeated_class_name": "MistralDecoderLayer" }
  ],

  // coarse dataflow. kinds: data | residual | mask | position | cross_attention | route | cache_read | cache_write
  "edges": [
    { "source": "embed_tokens", "target": "decoder_layers", "kind": "data" },
    { "source": "decoder_layers", "target": "norm", "kind": "data" },
    // block-level flow (pre-norm), re-grounded from the observed forward:
    { "source": "decoder_layer.input_layernorm", "target": "decoder_layer.self_attn", "kind": "data", "provenance": "observed_forward" },
    { "source": "decoder_layer.self_attn", "target": "decoder_layer.post_attention_layernorm", "kind": "data", "provenance": "observed_forward" },
    { "source": "decoder_layer.post_attention_layernorm", "target": "decoder_layer.mlp", "kind": "data", "provenance": "observed_forward" },
    // intra-module fan-out/fan-in (role-based; q/k/v parallel, not chained):
    { "source": "decoder_layer.self_attn", "target": "decoder_layer.self_attn.q_proj", "kind": "data", "provenance": "intra_module" },
    { "source": "decoder_layer.self_attn.q_proj", "target": "decoder_layer.self_attn.o_proj", "kind": "data", "provenance": "intra_module" },
    { "source": "decoder_layer.mlp.gate_proj", "target": "decoder_layer.mlp.down_proj", "kind": "data", "provenance": "intra_module" },
    { "source": "decoder_layer", "target": "decoder_layer.self_attn", "kind": "residual" },
    { "source": "rotary_emb", "target": "decoder_layer.self_attn", "kind": "position" },
    { "source": "input:attention_mask", "target": "decoder_layer.self_attn", "kind": "mask" },
    { "source": "state:kv_cache", "target": "decoder_layer.self_attn", "kind": "cache_read" },
    { "source": "decoder_layer.self_attn", "target": "state:kv_cache", "kind": "cache_write" }
  ],

  // observed tensor shapes (symbolized), keyed by node id; order lives in edges
  "dataflow": {
    "source": "observed_forward_meta",
    "input": { "name": "input_ids", "shape": ["B", "S"] },
    "output": { "shape": ["B", "S", "config.hidden_size"] },
    "shapes": {
      "embed_tokens": { "in": ["B", "S"], "out": ["B", "S", "config.hidden_size"] },
      "decoder_layer.self_attn": { "in": ["B", "S", "config.hidden_size"], "out": ["B", "S", "config.hidden_size"] }
      // … norm, mlp, the other block children
    }
  },

  // which classes this IR was introspected from (resolution strategy is invariant → in SPEC, not here)
  "provenance": {
    "config_class": "MistralConfig", "config_module": "transformers.models.mistral.configuration_mistral",
    "model_class": "MistralModel", "model_module": "transformers.models.mistral.modeling_mistral"
  }
}