{
  "architecture": {
    "attention_variant": "MHA",
    "family": "masked_lm",
    "is_moe": false,
    "mixer": "attention",
    "positional": "learned",
    "tie_word_embeddings": true,
    "view": "encoder"
  },
  "capabilities": {
    "attention_backends": [
      "eager",
      "sdpa",
      "flash_attention",
      "flex_attention"
    ],
    "attention_patterns": [
      "bidirectional"
    ],
    "attention_schedule": null,
    "task_heads": [
      "masked_lm",
      "multiple_choice",
      "pretraining",
      "question_answering",
      "sequence_classification",
      "text_encoding",
      "token_classification"
    ],
    "tensor_parallel": false
  },
  "components": [
    {
      "children": [
        "embeddings",
        "encoder",
        "pooler",
        "pooler_activation"
      ],
      "class_name": "AlbertModel",
      "id": "model",
      "kind": "model",
      "path_pattern": "model"
    },
    {
      "attributes": {
        "embedding_dim": "config.hidden_size",
        "num_embeddings": "config.vocab_size"
      },
      "children": [
        "embeddings.word_embeddings",
        "embeddings.position_embeddings",
        "embeddings.token_type_embeddings",
        "embeddings.LayerNorm"
      ],
      "class_name": "AlbertEmbeddings",
      "id": "embeddings",
      "kind": "embedding",
      "path_pattern": "model.embeddings"
    },
    {
      "attributes": {
        "embedding_dim": 128,
        "num_embeddings": "config.vocab_size",
        "tied_lm_head": true
      },
      "class_name": "Embedding",
      "id": "embeddings.word_embeddings",
      "kind": "embedding",
      "path_pattern": "model.embeddings.word_embeddings"
    },
    {
      "attributes": {
        "max_position_embeddings": "config.max_position_embeddings",
        "scheme": "learned"
      },
      "class_name": "Embedding",
      "id": "embeddings.position_embeddings",
      "kind": "position",
      "path_pattern": "model.embeddings.position_embeddings"
    },
    {
      "attributes": {
        "embedding_dim": 128,
        "num_embeddings": 2
      },
      "class_name": "Embedding",
      "id": "embeddings.token_type_embeddings",
      "kind": "embedding",
      "path_pattern": "model.embeddings.token_type_embeddings"
    },
    {
      "attributes": {
        "norm_type": "layer"
      },
      "class_name": "LayerNorm",
      "id": "embeddings.LayerNorm",
      "kind": "normalization",
      "path_pattern": "model.embeddings.LayerNorm"
    },
    {
      "children": [
        "encoder.albert_layer_groups"
      ],
      "class_name": "AlbertTransformer",
      "id": "encoder",
      "kind": "encoder",
      "path_pattern": "model.encoder"
    },
    {
      "children": [
        "encoder.albert_layer_groups.0"
      ],
      "class_name": "ModuleList",
      "id": "encoder.albert_layer_groups",
      "kind": "repeated_container",
      "path_pattern": "model.encoder.albert_layer_groups"
    },
    {
      "children": [
        "encoder.albert_layer_groups.0.albert_layers"
      ],
      "class_name": "AlbertLayerGroup",
      "id": "encoder.albert_layer_groups.0",
      "kind": "module",
      "path_pattern": "model.encoder.albert_layer_groups.0"
    },
    {
      "children": [
        "encoder.albert_layer_groups.0.albert_layers.0"
      ],
      "class_name": "ModuleList",
      "id": "encoder.albert_layer_groups.0.albert_layers",
      "kind": "repeated_container",
      "path_pattern": "model.encoder.albert_layer_groups.0.albert_layers"
    },
    {
      "children": [
        "encoder.albert_layer_groups.0.albert_layers.0.full_layer_layer_norm",
        "encoder.albert_layer_groups.0.albert_layers.0.attention"
      ],
      "class_name": "AlbertLayer",
      "id": "encoder.albert_layer_groups.0.albert_layers.0",
      "kind": "transformer_block",
      "path_pattern": "model.encoder.albert_layer_groups.0.albert_layers.0"
    },
    {
      "attributes": {
        "norm_type": "layer"
      },
      "class_name": "LayerNorm",
      "id": "encoder.albert_layer_groups.0.albert_layers.0.full_layer_layer_norm",
      "kind": "normalization",
      "path_pattern": "model.encoder.albert_layer_groups.0.albert_layers.0.full_layer_layer_norm"
    },
    {
      "attributes": {
        "head_dim": 64,
        "n_heads": 64,
        "n_kv_heads": 64,
        "pattern": "bidirectional",
        "rope": false,
        "variant": "MHA"
      },
      "children": [
        "encoder.albert_layer_groups.0.albert_layers.0.attention.query",
        "encoder.albert_layer_groups.0.albert_layers.0.attention.key",
        "encoder.albert_layer_groups.0.albert_layers.0.attention.value",
        "encoder.albert_layer_groups.0.albert_layers.0.attention.dense",
        "encoder.albert_layer_groups.0.albert_layers.0.attention.LayerNorm"
      ],
      "class_name": "AlbertAttention",
      "id": "encoder.albert_layer_groups.0.albert_layers.0.attention",
      "kind": "attention",
      "path_pattern": "model.encoder.albert_layer_groups.0.albert_layers.0.attention"
    },
    {
      "attributes": {
        "in_features": "config.hidden_size",
        "out_features": "config.hidden_size"
      },
      "class_name": "Linear",
      "id": "encoder.albert_layer_groups.0.albert_layers.0.attention.query",
      "kind": "projection",
      "path_pattern": "model.encoder.albert_layer_groups.0.albert_layers.0.attention.query"
    },
    {
      "attributes": {
        "in_features": "config.hidden_size",
        "out_features": "config.hidden_size"
      },
      "class_name": "Linear",
      "id": "encoder.albert_layer_groups.0.albert_layers.0.attention.key",
      "kind": "projection",
      "path_pattern": "model.encoder.albert_layer_groups.0.albert_layers.0.attention.key"
    },
    {
      "attributes": {
        "in_features": "config.hidden_size",
        "out_features": "config.hidden_size"
      },
      "class_name": "Linear",
      "id": "encoder.albert_layer_groups.0.albert_layers.0.attention.value",
      "kind": "projection",
      "path_pattern": "model.encoder.albert_layer_groups.0.albert_layers.0.attention.value"
    },
    {
      "attributes": {
        "in_features": "config.hidden_size",
        "out_features": "config.hidden_size"
      },
      "class_name": "Linear",
      "id": "encoder.albert_layer_groups.0.albert_layers.0.attention.dense",
      "kind": "projection",
      "path_pattern": "model.encoder.albert_layer_groups.0.albert_layers.0.attention.dense"
    },
    {
      "attributes": {
        "norm_type": "layer"
      },
      "class_name": "LayerNorm",
      "id": "encoder.albert_layer_groups.0.albert_layers.0.attention.LayerNorm",
      "kind": "normalization",
      "path_pattern": "model.encoder.albert_layer_groups.0.albert_layers.0.attention.LayerNorm"
    },
    {
      "attributes": {
        "in_features": "config.hidden_size",
        "out_features": "config.hidden_size"
      },
      "class_name": "Linear",
      "id": "pooler",
      "kind": "projection",
      "path_pattern": "model.pooler"
    },
    {
      "class_name": "Tanh",
      "id": "pooler_activation",
      "kind": "module",
      "path_pattern": "model.pooler_activation"
    }
  ],
  "config": {
    "class_name": "AlbertConfig",
    "model_type": "albert",
    "module": "transformers.models.albert.configuration_albert",
    "referenced_fields": {},
    "salient_fields": {
      "hidden_act": "gelu_new",
      "hidden_size": 4096,
      "intermediate_size": 16384,
      "is_encoder_decoder": false,
      "max_position_embeddings": 512,
      "num_attention_heads": 64,
      "num_hidden_layers": 12,
      "tie_word_embeddings": true,
      "vocab_size": 30000
    }
  },
  "dataflow": {
    "input": {
      "name": "input_ids",
      "shape": [
        "B",
        "S"
      ]
    },
    "output": {
      "shape": [
        "B",
        "S",
        "config.hidden_size"
      ]
    },
    "shapes": {
      "embeddings": {
        "in": [
          "B",
          "S"
        ],
        "out": [
          "B",
          "S",
          128
        ]
      },
      "encoder": {
        "in": [
          "B",
          "S",
          128
        ],
        "out": [
          "B",
          "S",
          "config.hidden_size"
        ]
      },
      "pooler": {
        "in": [
          "B",
          "config.hidden_size"
        ],
        "out": [
          "B",
          "config.hidden_size"
        ]
      },
      "pooler_activation": {
        "in": [
          "B",
          "config.hidden_size"
        ],
        "out": [
          "B",
          "config.hidden_size"
        ]
      }
    },
    "source": "observed_forward_meta"
  },
  "edges": [
    {
      "kind": "data",
      "source": "embeddings",
      "target": "encoder"
    },
    {
      "kind": "data",
      "source": "embeddings.word_embeddings",
      "target": "embeddings.token_type_embeddings"
    },
    {
      "kind": "data",
      "source": "embeddings.token_type_embeddings",
      "target": "embeddings.LayerNorm"
    },
    {
      "kind": "data",
      "source": "encoder.albert_layer_groups.0.albert_layers.0.full_layer_layer_norm",
      "target": "encoder.albert_layer_groups.0.albert_layers.0.attention"
    },
    {
      "kind": "mask",
      "source": "input:attention_mask",
      "target": "encoder.albert_layer_groups.0.albert_layers.0.attention"
    },
    {
      "kind": "residual",
      "source": "encoder.albert_layer_groups.0.albert_layers.0",
      "target": "encoder.albert_layer_groups.0.albert_layers.0.attention"
    },
    {
      "kind": "position",
      "source": "embeddings.position_embeddings",
      "target": "embeddings"
    },
    {
      "kind": "data",
      "provenance": "intra_module",
      "source": "encoder.albert_layer_groups.0.albert_layers.0.attention",
      "target": "encoder.albert_layer_groups.0.albert_layers.0.attention.query"
    },
    {
      "kind": "data",
      "provenance": "intra_module",
      "source": "encoder.albert_layer_groups.0.albert_layers.0.attention",
      "target": "encoder.albert_layer_groups.0.albert_layers.0.attention.key"
    },
    {
      "kind": "data",
      "provenance": "intra_module",
      "source": "encoder.albert_layer_groups.0.albert_layers.0.attention",
      "target": "encoder.albert_layer_groups.0.albert_layers.0.attention.value"
    },
    {
      "kind": "data",
      "provenance": "intra_module",
      "source": "encoder.albert_layer_groups.0.albert_layers.0.attention",
      "target": "encoder.albert_layer_groups.0.albert_layers.0.attention.dense"
    },
    {
      "kind": "data",
      "provenance": "observed_forward",
      "source": "encoder",
      "target": "pooler"
    },
    {
      "kind": "data",
      "provenance": "observed_forward",
      "source": "pooler",
      "target": "pooler_activation"
    }
  ],
  "extends": null,
  "model_type": "albert",
  "provenance": {
    "config_class": "AlbertConfig",
    "config_module": "transformers.models.albert.configuration_albert",
    "model_class": "AlbertModel",
    "model_module": "transformers.models.albert.modeling_albert"
  },
  "repeats": [],
  "schema_version": "architecture-template-v0",
  "templates": []
}