{
  "architecture": {
    "attention_variant": "MHA",
    "family": "causal_lm",
    "is_moe": false,
    "mixer": "attention",
    "tie_word_embeddings": true,
    "view": "decoder"
  },
  "capabilities": {
    "attention_backends": [
      "eager"
    ],
    "attention_patterns": [
      "causal"
    ],
    "attention_schedule": null,
    "task_heads": [
      "causal_lm",
      "pretraining",
      "question_answering",
      "sequence_classification",
      "token_classification"
    ],
    "tensor_parallel": false
  },
  "components": [
    {
      "children": [
        "word_embeddings",
        "word_embeddings_layernorm",
        "blocks",
        "ln_f"
      ],
      "class_name": "BloomModel",
      "id": "model",
      "kind": "model",
      "path_pattern": "model"
    },
    {
      "attributes": {
        "embedding_dim": "config.hidden_size",
        "num_embeddings": "config.vocab_size",
        "tied_lm_head": true
      },
      "class_name": "Embedding",
      "id": "word_embeddings",
      "kind": "embedding",
      "path_pattern": "model.word_embeddings"
    },
    {
      "attributes": {
        "norm_type": "layer"
      },
      "class_name": "LayerNorm",
      "id": "word_embeddings_layernorm",
      "kind": "normalization",
      "path_pattern": "model.word_embeddings_layernorm"
    },
    {
      "attributes": {
        "norm_type": "layer"
      },
      "class_name": "LayerNorm",
      "id": "ln_f",
      "kind": "normalization",
      "path_pattern": "model.ln_f"
    }
  ],
  "config": {
    "class_name": "BloomConfig",
    "model_type": "bloom",
    "module": "transformers.models.bloom.configuration_bloom",
    "referenced_fields": {},
    "salient_fields": {
      "hidden_size": 64,
      "is_encoder_decoder": false,
      "num_attention_heads": 8,
      "num_hidden_layers": 2,
      "tie_word_embeddings": true,
      "vocab_size": 250880
    }
  },
  "dataflow": {
    "input": {
      "name": "input_ids",
      "shape": [
        "B",
        "S"
      ]
    },
    "output": {
      "shape": [
        "B",
        "S",
        "config.hidden_size"
      ]
    },
    "shapes": {
      "block.input_layernorm": {
        "in": [
          "B",
          "S",
          "config.hidden_size"
        ],
        "out": [
          "B",
          "S",
          "config.hidden_size"
        ]
      },
      "block.mlp": {
        "in": [
          "B",
          "S",
          "config.hidden_size"
        ],
        "out": [
          "B",
          "S",
          "config.hidden_size"
        ]
      },
      "block.post_attention_layernorm": {
        "in": [
          "B",
          "S",
          "config.hidden_size"
        ],
        "out": [
          "B",
          "S",
          "config.hidden_size"
        ]
      },
      "block.self_attention": {
        "in": [
          "B",
          "S",
          "config.hidden_size"
        ],
        "out": [
          "B",
          "S",
          "config.hidden_size"
        ]
      },
      "blocks": {
        "in": [
          "B",
          "S",
          "config.hidden_size"
        ],
        "out": [
          "B",
          "S",
          "config.hidden_size"
        ]
      },
      "ln_f": {
        "in": [
          "B",
          "S",
          "config.hidden_size"
        ],
        "out": [
          "B",
          "S",
          "config.hidden_size"
        ]
      },
      "word_embeddings": {
        "in": [
          "B",
          "S"
        ],
        "out": [
          "B",
          "S",
          "config.hidden_size"
        ]
      },
      "word_embeddings_layernorm": {
        "in": [
          "B",
          "S",
          "config.hidden_size"
        ],
        "out": [
          "B",
          "S",
          "config.hidden_size"
        ]
      }
    },
    "source": "observed_forward_meta"
  },
  "edges": [
    {
      "kind": "data",
      "source": "word_embeddings",
      "target": "word_embeddings_layernorm"
    },
    {
      "kind": "data",
      "source": "word_embeddings_layernorm",
      "target": "blocks"
    },
    {
      "kind": "data",
      "source": "blocks",
      "target": "ln_f"
    },
    {
      "kind": "mask",
      "source": "input:attention_mask",
      "target": "block.self_attention"
    },
    {
      "kind": "residual",
      "source": "block",
      "target": "block.self_attention"
    },
    {
      "kind": "residual",
      "source": "block",
      "target": "block.mlp"
    },
    {
      "kind": "data",
      "provenance": "intra_module",
      "source": "block.self_attention",
      "target": "block.self_attention.query_key_value"
    },
    {
      "kind": "data",
      "provenance": "intra_module",
      "source": "block.self_attention",
      "target": "block.self_attention.dense"
    },
    {
      "kind": "data",
      "provenance": "intra_module",
      "source": "block.mlp",
      "target": "block.mlp.dense_h_to_4h"
    },
    {
      "kind": "data",
      "provenance": "intra_module",
      "source": "block.mlp.dense_h_to_4h",
      "target": "block.mlp.dense_4h_to_h"
    },
    {
      "kind": "cache_read",
      "source": "state:kv_cache",
      "target": "block.self_attention"
    },
    {
      "kind": "cache_write",
      "source": "block.self_attention",
      "target": "state:kv_cache"
    },
    {
      "kind": "data",
      "provenance": "observed_forward",
      "source": "block.input_layernorm",
      "target": "block.self_attention"
    },
    {
      "kind": "data",
      "provenance": "observed_forward",
      "source": "block.self_attention",
      "target": "block.post_attention_layernorm"
    },
    {
      "kind": "data",
      "provenance": "observed_forward",
      "source": "block.post_attention_layernorm",
      "target": "block.mlp"
    }
  ],
  "extends": null,
  "model_type": "bloom",
  "provenance": {
    "config_class": "BloomConfig",
    "config_module": "transformers.models.bloom.configuration_bloom",
    "model_class": "BloomModel",
    "model_module": "transformers.models.bloom.modeling_bloom"
  },
  "repeats": [
    {
      "body": "block",
      "container_path_pattern": "model.h",
      "count": 2,
      "count_expr": "config.num_hidden_layers",
      "count_source": "config",
      "id": "blocks",
      "index_symbol": "i",
      "item_path_pattern": "model.h.{i}",
      "kind": "symbolic_repeat",
      "provenance": {
        "class_name": "BloomBlock",
        "container_path_pattern": "model.h",
        "item_path_pattern": "model.h.{i}",
        "source": "module_tree_repeat_collapse"
      },
      "repeated_class_name": "BloomBlock"
    }
  ],
  "schema_version": "architecture-template-v0",
  "templates": [
    {
      "children": [
        "block.input_layernorm",
        "block.self_attention",
        "block.post_attention_layernorm",
        "block.mlp"
      ],
      "class_name": "BloomBlock",
      "id": "block",
      "kind": "transformer_block",
      "path_pattern": "model.h.{i}"
    },
    {
      "attributes": {
        "norm_type": "layer"
      },
      "class_name": "LayerNorm",
      "id": "block.input_layernorm",
      "kind": "normalization",
      "path_pattern": "model.h.{i}.input_layernorm"
    },
    {
      "attributes": {
        "head_dim": 8,
        "n_heads": 8,
        "n_kv_heads": 8,
        "pattern": "causal",
        "rope": false,
        "variant": "MHA"
      },
      "children": [
        "block.self_attention.query_key_value",
        "block.self_attention.dense"
      ],
      "class_name": "BloomAttention",
      "id": "block.self_attention",
      "kind": "attention",
      "path_pattern": "model.h.{i}.self_attention"
    },
    {
      "attributes": {
        "in_features": "config.hidden_size",
        "out_features": 192
      },
      "class_name": "Linear",
      "id": "block.self_attention.query_key_value",
      "kind": "projection",
      "path_pattern": "model.h.{i}.self_attention.query_key_value"
    },
    {
      "attributes": {
        "in_features": "config.hidden_size",
        "out_features": "config.hidden_size"
      },
      "class_name": "Linear",
      "id": "block.self_attention.dense",
      "kind": "projection",
      "path_pattern": "model.h.{i}.self_attention.dense"
    },
    {
      "attributes": {
        "norm_type": "layer"
      },
      "class_name": "LayerNorm",
      "id": "block.post_attention_layernorm",
      "kind": "normalization",
      "path_pattern": "model.h.{i}.post_attention_layernorm"
    },
    {
      "attributes": {
        "hidden_size": 64
      },
      "children": [
        "block.mlp.dense_h_to_4h",
        "block.mlp.gelu_impl",
        "block.mlp.dense_4h_to_h"
      ],
      "class_name": "BloomMLP",
      "id": "block.mlp",
      "kind": "feed_forward",
      "path_pattern": "model.h.{i}.mlp"
    },
    {
      "attributes": {
        "in_features": "config.hidden_size",
        "out_features": 256
      },
      "class_name": "Linear",
      "id": "block.mlp.dense_h_to_4h",
      "kind": "projection",
      "path_pattern": "model.h.{i}.mlp.dense_h_to_4h"
    },
    {
      "attributes": {
        "hidden_size": 64
      },
      "class_name": "BloomGelu",
      "id": "block.mlp.gelu_impl",
      "kind": "feed_forward",
      "path_pattern": "model.h.{i}.mlp.gelu_impl"
    },
    {
      "attributes": {
        "in_features": 256,
        "out_features": "config.hidden_size"
      },
      "class_name": "Linear",
      "id": "block.mlp.dense_4h_to_h",
      "kind": "projection",
      "path_pattern": "model.h.{i}.mlp.dense_4h_to_h"
    }
  ]
}