{
  "id": "jamba",
  "name": "Jamba",
  "description": "Jamba (AI21 Labs): full graph; dimensions verified against the official config.json.",
  "components": [
    {
      "id": "n1",
      "type": "input",
      "name": "tokens",
      "position": {
        "x": 300,
        "y": 50
      },
      "params": {
        "shape": [
          1,
          4096
        ]
      },
      "inputs": [],
      "outputs": [
        "n2"
      ]
    },
    {
      "id": "n2",
      "type": "embedding",
      "name": "token_embed",
      "position": {
        "x": 300,
        "y": 180
      },
      "params": {
        "numEmbeddings": 65536,
        "embeddingDim": 4096
      },
      "inputs": [
        "n1"
      ],
      "outputs": [
        "n3",
        "n5"
      ],
      "scope": "embeddings"
    },
    {
      "id": "n3",
      "type": "rmsNorm",
      "name": "mixer_norm_1",
      "position": {
        "x": 300,
        "y": 310
      },
      "params": {
        "normalizedShape": 4096
      },
      "inputs": [
        "n2"
      ],
      "outputs": [
        "n4"
      ],
      "scope": "layer.0",
      "notes": "One 8-layer Jamba block: 7 Mamba mixers + 1 attention (layer 5), MoE every other layer. The full model stacks 4 of these (32 layers)."
    },
    {
      "id": "n4",
      "type": "mamba",
      "name": "mamba_1",
      "position": {
        "x": 300,
        "y": 440
      },
      "params": {
        "dModel": 4096,
        "dState": 16,
        "expand": 2
      },
      "inputs": [
        "n3"
      ],
      "outputs": [
        "n5"
      ],
      "scope": "layer.0"
    },
    {
      "id": "n5",
      "type": "add",
      "name": "mixer_residual_1",
      "position": {
        "x": 300,
        "y": 570
      },
      "params": {},
      "inputs": [
        "n2",
        "n4"
      ],
      "outputs": [
        "n6",
        "n8"
      ],
      "scope": "layer.0"
    },
    {
      "id": "n6",
      "type": "rmsNorm",
      "name": "ffn_norm_1",
      "position": {
        "x": 300,
        "y": 700
      },
      "params": {
        "normalizedShape": 4096
      },
      "inputs": [
        "n5"
      ],
      "outputs": [
        "n7"
      ],
      "scope": "layer.0"
    },
    {
      "id": "n7",
      "type": "feedForward",
      "name": "mlp_1",
      "position": {
        "x": 300,
        "y": 830
      },
      "params": {
        "embedDim": 4096,
        "ffDim": 14336
      },
      "inputs": [
        "n6"
      ],
      "outputs": [
        "n8"
      ],
      "scope": "layer.0"
    },
    {
      "id": "n8",
      "type": "add",
      "name": "ffn_residual_1",
      "position": {
        "x": 300,
        "y": 960
      },
      "params": {},
      "inputs": [
        "n5",
        "n7"
      ],
      "outputs": [
        "n9",
        "n11"
      ],
      "scope": "layer.0"
    },
    {
      "id": "n9",
      "type": "rmsNorm",
      "name": "mixer_norm_2",
      "position": {
        "x": 300,
        "y": 1090
      },
      "params": {
        "normalizedShape": 4096
      },
      "inputs": [
        "n8"
      ],
      "outputs": [
        "n10"
      ],
      "scope": "layer.1"
    },
    {
      "id": "n10",
      "type": "mamba",
      "name": "mamba_2",
      "position": {
        "x": 300,
        "y": 1220
      },
      "params": {
        "dModel": 4096,
        "dState": 16,
        "expand": 2
      },
      "inputs": [
        "n9"
      ],
      "outputs": [
        "n11"
      ],
      "scope": "layer.1"
    },
    {
      "id": "n11",
      "type": "add",
      "name": "mixer_residual_2",
      "position": {
        "x": 300,
        "y": 1350
      },
      "params": {},
      "inputs": [
        "n8",
        "n10"
      ],
      "outputs": [
        "n12",
        "n14"
      ],
      "scope": "layer.1"
    },
    {
      "id": "n12",
      "type": "rmsNorm",
      "name": "ffn_norm_2",
      "position": {
        "x": 300,
        "y": 1480
      },
      "params": {
        "normalizedShape": 4096
      },
      "inputs": [
        "n11"
      ],
      "outputs": [
        "n13"
      ],
      "scope": "layer.1"
    },
    {
      "id": "n13",
      "type": "moeLayer",
      "name": "moe_2",
      "position": {
        "x": 300,
        "y": 1610
      },
      "params": {
        "embedDim": 4096,
        "numExperts": 16,
        "topK": 2,
        "expertDim": 14336
      },
      "inputs": [
        "n12"
      ],
      "outputs": [
        "n14"
      ],
      "scope": "layer.1"
    },
    {
      "id": "n14",
      "type": "add",
      "name": "ffn_residual_2",
      "position": {
        "x": 300,
        "y": 1740
      },
      "params": {},
      "inputs": [
        "n11",
        "n13"
      ],
      "outputs": [
        "n15",
        "n17"
      ],
      "scope": "layer.1"
    },
    {
      "id": "n15",
      "type": "rmsNorm",
      "name": "mixer_norm_3",
      "position": {
        "x": 300,
        "y": 1870
      },
      "params": {
        "normalizedShape": 4096
      },
      "inputs": [
        "n14"
      ],
      "outputs": [
        "n16"
      ],
      "scope": "layer.2"
    },
    {
      "id": "n16",
      "type": "mamba",
      "name": "mamba_3",
      "position": {
        "x": 300,
        "y": 2000
      },
      "params": {
        "dModel": 4096,
        "dState": 16,
        "expand": 2
      },
      "inputs": [
        "n15"
      ],
      "outputs": [
        "n17"
      ],
      "scope": "layer.2"
    },
    {
      "id": "n17",
      "type": "add",
      "name": "mixer_residual_3",
      "position": {
        "x": 300,
        "y": 2130
      },
      "params": {},
      "inputs": [
        "n14",
        "n16"
      ],
      "outputs": [
        "n18",
        "n20"
      ],
      "scope": "layer.2"
    },
    {
      "id": "n18",
      "type": "rmsNorm",
      "name": "ffn_norm_3",
      "position": {
        "x": 300,
        "y": 2260
      },
      "params": {
        "normalizedShape": 4096
      },
      "inputs": [
        "n17"
      ],
      "outputs": [
        "n19"
      ],
      "scope": "layer.2"
    },
    {
      "id": "n19",
      "type": "feedForward",
      "name": "mlp_3",
      "position": {
        "x": 300,
        "y": 2390
      },
      "params": {
        "embedDim": 4096,
        "ffDim": 14336
      },
      "inputs": [
        "n18"
      ],
      "outputs": [
        "n20"
      ],
      "scope": "layer.2"
    },
    {
      "id": "n20",
      "type": "add",
      "name": "ffn_residual_3",
      "position": {
        "x": 300,
        "y": 2520
      },
      "params": {},
      "inputs": [
        "n17",
        "n19"
      ],
      "outputs": [
        "n21",
        "n23"
      ],
      "scope": "layer.2"
    },
    {
      "id": "n21",
      "type": "rmsNorm",
      "name": "mixer_norm_4",
      "position": {
        "x": 300,
        "y": 2650
      },
      "params": {
        "normalizedShape": 4096
      },
      "inputs": [
        "n20"
      ],
      "outputs": [
        "n22"
      ],
      "scope": "layer.3"
    },
    {
      "id": "n22",
      "type": "mamba",
      "name": "mamba_4",
      "position": {
        "x": 300,
        "y": 2780
      },
      "params": {
        "dModel": 4096,
        "dState": 16,
        "expand": 2
      },
      "inputs": [
        "n21"
      ],
      "outputs": [
        "n23"
      ],
      "scope": "layer.3"
    },
    {
      "id": "n23",
      "type": "add",
      "name": "mixer_residual_4",
      "position": {
        "x": 300,
        "y": 2910
      },
      "params": {},
      "inputs": [
        "n20",
        "n22"
      ],
      "outputs": [
        "n24",
        "n26"
      ],
      "scope": "layer.3"
    },
    {
      "id": "n24",
      "type": "rmsNorm",
      "name": "ffn_norm_4",
      "position": {
        "x": 300,
        "y": 3040
      },
      "params": {
        "normalizedShape": 4096
      },
      "inputs": [
        "n23"
      ],
      "outputs": [
        "n25"
      ],
      "scope": "layer.3"
    },
    {
      "id": "n25",
      "type": "moeLayer",
      "name": "moe_4",
      "position": {
        "x": 300,
        "y": 3170
      },
      "params": {
        "embedDim": 4096,
        "numExperts": 16,
        "topK": 2,
        "expertDim": 14336
      },
      "inputs": [
        "n24"
      ],
      "outputs": [
        "n26"
      ],
      "scope": "layer.3"
    },
    {
      "id": "n26",
      "type": "add",
      "name": "ffn_residual_4",
      "position": {
        "x": 300,
        "y": 3300
      },
      "params": {},
      "inputs": [
        "n23",
        "n25"
      ],
      "outputs": [
        "n27",
        "n29"
      ],
      "scope": "layer.3"
    },
    {
      "id": "n27",
      "type": "rmsNorm",
      "name": "mixer_norm_5",
      "position": {
        "x": 300,
        "y": 3430
      },
      "params": {
        "normalizedShape": 4096
      },
      "inputs": [
        "n26"
      ],
      "outputs": [
        "n28"
      ],
      "scope": "layer.4"
    },
    {
      "id": "n28",
      "type": "groupedQueryAttention",
      "name": "attn_5",
      "position": {
        "x": 300,
        "y": 3560
      },
      "params": {
        "embedDim": 4096,
        "numHeads": 32,
        "numKVHeads": 8,
        "headDim": 128
      },
      "inputs": [
        "n27"
      ],
      "outputs": [
        "n29"
      ],
      "scope": "layer.4"
    },
    {
      "id": "n29",
      "type": "add",
      "name": "mixer_residual_5",
      "position": {
        "x": 300,
        "y": 3690
      },
      "params": {},
      "inputs": [
        "n26",
        "n28"
      ],
      "outputs": [
        "n30",
        "n32"
      ],
      "scope": "layer.4"
    },
    {
      "id": "n30",
      "type": "rmsNorm",
      "name": "ffn_norm_5",
      "position": {
        "x": 300,
        "y": 3820
      },
      "params": {
        "normalizedShape": 4096
      },
      "inputs": [
        "n29"
      ],
      "outputs": [
        "n31"
      ],
      "scope": "layer.4"
    },
    {
      "id": "n31",
      "type": "feedForward",
      "name": "mlp_5",
      "position": {
        "x": 300,
        "y": 3950
      },
      "params": {
        "embedDim": 4096,
        "ffDim": 14336
      },
      "inputs": [
        "n30"
      ],
      "outputs": [
        "n32"
      ],
      "scope": "layer.4"
    },
    {
      "id": "n32",
      "type": "add",
      "name": "ffn_residual_5",
      "position": {
        "x": 300,
        "y": 4080
      },
      "params": {},
      "inputs": [
        "n29",
        "n31"
      ],
      "outputs": [
        "n33",
        "n35"
      ],
      "scope": "layer.4"
    },
    {
      "id": "n33",
      "type": "rmsNorm",
      "name": "mixer_norm_6",
      "position": {
        "x": 300,
        "y": 4210
      },
      "params": {
        "normalizedShape": 4096
      },
      "inputs": [
        "n32"
      ],
      "outputs": [
        "n34"
      ],
      "scope": "layer.5"
    },
    {
      "id": "n34",
      "type": "mamba",
      "name": "mamba_6",
      "position": {
        "x": 300,
        "y": 4340
      },
      "params": {
        "dModel": 4096,
        "dState": 16,
        "expand": 2
      },
      "inputs": [
        "n33"
      ],
      "outputs": [
        "n35"
      ],
      "scope": "layer.5"
    },
    {
      "id": "n35",
      "type": "add",
      "name": "mixer_residual_6",
      "position": {
        "x": 300,
        "y": 4470
      },
      "params": {},
      "inputs": [
        "n32",
        "n34"
      ],
      "outputs": [
        "n36",
        "n38"
      ],
      "scope": "layer.5"
    },
    {
      "id": "n36",
      "type": "rmsNorm",
      "name": "ffn_norm_6",
      "position": {
        "x": 300,
        "y": 4600
      },
      "params": {
        "normalizedShape": 4096
      },
      "inputs": [
        "n35"
      ],
      "outputs": [
        "n37"
      ],
      "scope": "layer.5"
    },
    {
      "id": "n37",
      "type": "moeLayer",
      "name": "moe_6",
      "position": {
        "x": 300,
        "y": 4730
      },
      "params": {
        "embedDim": 4096,
        "numExperts": 16,
        "topK": 2,
        "expertDim": 14336
      },
      "inputs": [
        "n36"
      ],
      "outputs": [
        "n38"
      ],
      "scope": "layer.5"
    },
    {
      "id": "n38",
      "type": "add",
      "name": "ffn_residual_6",
      "position": {
        "x": 300,
        "y": 4860
      },
      "params": {},
      "inputs": [
        "n35",
        "n37"
      ],
      "outputs": [
        "n39",
        "n41"
      ],
      "scope": "layer.5"
    },
    {
      "id": "n39",
      "type": "rmsNorm",
      "name": "mixer_norm_7",
      "position": {
        "x": 300,
        "y": 4990
      },
      "params": {
        "normalizedShape": 4096
      },
      "inputs": [
        "n38"
      ],
      "outputs": [
        "n40"
      ],
      "scope": "layer.6"
    },
    {
      "id": "n40",
      "type": "mamba",
      "name": "mamba_7",
      "position": {
        "x": 300,
        "y": 5120
      },
      "params": {
        "dModel": 4096,
        "dState": 16,
        "expand": 2
      },
      "inputs": [
        "n39"
      ],
      "outputs": [
        "n41"
      ],
      "scope": "layer.6"
    },
    {
      "id": "n41",
      "type": "add",
      "name": "mixer_residual_7",
      "position": {
        "x": 300,
        "y": 5250
      },
      "params": {},
      "inputs": [
        "n38",
        "n40"
      ],
      "outputs": [
        "n42",
        "n44"
      ],
      "scope": "layer.6"
    },
    {
      "id": "n42",
      "type": "rmsNorm",
      "name": "ffn_norm_7",
      "position": {
        "x": 300,
        "y": 5380
      },
      "params": {
        "normalizedShape": 4096
      },
      "inputs": [
        "n41"
      ],
      "outputs": [
        "n43"
      ],
      "scope": "layer.6"
    },
    {
      "id": "n43",
      "type": "feedForward",
      "name": "mlp_7",
      "position": {
        "x": 300,
        "y": 5510
      },
      "params": {
        "embedDim": 4096,
        "ffDim": 14336
      },
      "inputs": [
        "n42"
      ],
      "outputs": [
        "n44"
      ],
      "scope": "layer.6"
    },
    {
      "id": "n44",
      "type": "add",
      "name": "ffn_residual_7",
      "position": {
        "x": 300,
        "y": 5640
      },
      "params": {},
      "inputs": [
        "n41",
        "n43"
      ],
      "outputs": [
        "n45",
        "n47"
      ],
      "scope": "layer.6"
    },
    {
      "id": "n45",
      "type": "rmsNorm",
      "name": "mixer_norm_8",
      "position": {
        "x": 300,
        "y": 5770
      },
      "params": {
        "normalizedShape": 4096
      },
      "inputs": [
        "n44"
      ],
      "outputs": [
        "n46"
      ],
      "scope": "layer.7"
    },
    {
      "id": "n46",
      "type": "mamba",
      "name": "mamba_8",
      "position": {
        "x": 300,
        "y": 5900
      },
      "params": {
        "dModel": 4096,
        "dState": 16,
        "expand": 2
      },
      "inputs": [
        "n45"
      ],
      "outputs": [
        "n47"
      ],
      "scope": "layer.7"
    },
    {
      "id": "n47",
      "type": "add",
      "name": "mixer_residual_8",
      "position": {
        "x": 300,
        "y": 6030
      },
      "params": {},
      "inputs": [
        "n44",
        "n46"
      ],
      "outputs": [
        "n48",
        "n50"
      ],
      "scope": "layer.7"
    },
    {
      "id": "n48",
      "type": "rmsNorm",
      "name": "ffn_norm_8",
      "position": {
        "x": 300,
        "y": 6160
      },
      "params": {
        "normalizedShape": 4096
      },
      "inputs": [
        "n47"
      ],
      "outputs": [
        "n49"
      ],
      "scope": "layer.7"
    },
    {
      "id": "n49",
      "type": "moeLayer",
      "name": "moe_8",
      "position": {
        "x": 300,
        "y": 6290
      },
      "params": {
        "embedDim": 4096,
        "numExperts": 16,
        "topK": 2,
        "expertDim": 14336
      },
      "inputs": [
        "n48"
      ],
      "outputs": [
        "n50"
      ],
      "scope": "layer.7"
    },
    {
      "id": "n50",
      "type": "add",
      "name": "ffn_residual_8",
      "position": {
        "x": 300,
        "y": 6420
      },
      "params": {},
      "inputs": [
        "n47",
        "n49"
      ],
      "outputs": [
        "n51"
      ],
      "scope": "layer.7"
    },
    {
      "id": "n51",
      "type": "rmsNorm",
      "name": "final_norm",
      "position": {
        "x": 300,
        "y": 6550
      },
      "params": {
        "normalizedShape": 4096
      },
      "inputs": [
        "n50"
      ],
      "outputs": [
        "n52"
      ],
      "scope": "norm"
    },
    {
      "id": "n52",
      "type": "linear",
      "name": "lm_head",
      "position": {
        "x": 300,
        "y": 6680
      },
      "params": {
        "inFeatures": 4096,
        "outFeatures": 65536
      },
      "inputs": [
        "n51"
      ],
      "outputs": [
        "n53"
      ],
      "scope": "head"
    },
    {
      "id": "n53",
      "type": "output",
      "name": "logits",
      "position": {
        "x": 300,
        "y": 6810
      },
      "params": {},
      "inputs": [
        "n52"
      ],
      "outputs": []
    }
  ],
  "connections": [
    {
      "id": "c1",
      "from": "n1",
      "to": "n2"
    },
    {
      "id": "c2",
      "from": "n2",
      "to": "n3"
    },
    {
      "id": "c3",
      "from": "n3",
      "to": "n4"
    },
    {
      "id": "c4",
      "from": "n2",
      "to": "n5"
    },
    {
      "id": "c5",
      "from": "n4",
      "to": "n5"
    },
    {
      "id": "c6",
      "from": "n5",
      "to": "n6"
    },
    {
      "id": "c7",
      "from": "n6",
      "to": "n7"
    },
    {
      "id": "c8",
      "from": "n5",
      "to": "n8"
    },
    {
      "id": "c9",
      "from": "n7",
      "to": "n8"
    },
    {
      "id": "c10",
      "from": "n8",
      "to": "n9"
    },
    {
      "id": "c11",
      "from": "n9",
      "to": "n10"
    },
    {
      "id": "c12",
      "from": "n8",
      "to": "n11"
    },
    {
      "id": "c13",
      "from": "n10",
      "to": "n11"
    },
    {
      "id": "c14",
      "from": "n11",
      "to": "n12"
    },
    {
      "id": "c15",
      "from": "n12",
      "to": "n13"
    },
    {
      "id": "c16",
      "from": "n11",
      "to": "n14"
    },
    {
      "id": "c17",
      "from": "n13",
      "to": "n14"
    },
    {
      "id": "c18",
      "from": "n14",
      "to": "n15"
    },
    {
      "id": "c19",
      "from": "n15",
      "to": "n16"
    },
    {
      "id": "c20",
      "from": "n14",
      "to": "n17"
    },
    {
      "id": "c21",
      "from": "n16",
      "to": "n17"
    },
    {
      "id": "c22",
      "from": "n17",
      "to": "n18"
    },
    {
      "id": "c23",
      "from": "n18",
      "to": "n19"
    },
    {
      "id": "c24",
      "from": "n17",
      "to": "n20"
    },
    {
      "id": "c25",
      "from": "n19",
      "to": "n20"
    },
    {
      "id": "c26",
      "from": "n20",
      "to": "n21"
    },
    {
      "id": "c27",
      "from": "n21",
      "to": "n22"
    },
    {
      "id": "c28",
      "from": "n20",
      "to": "n23"
    },
    {
      "id": "c29",
      "from": "n22",
      "to": "n23"
    },
    {
      "id": "c30",
      "from": "n23",
      "to": "n24"
    },
    {
      "id": "c31",
      "from": "n24",
      "to": "n25"
    },
    {
      "id": "c32",
      "from": "n23",
      "to": "n26"
    },
    {
      "id": "c33",
      "from": "n25",
      "to": "n26"
    },
    {
      "id": "c34",
      "from": "n26",
      "to": "n27"
    },
    {
      "id": "c35",
      "from": "n27",
      "to": "n28"
    },
    {
      "id": "c36",
      "from": "n26",
      "to": "n29"
    },
    {
      "id": "c37",
      "from": "n28",
      "to": "n29"
    },
    {
      "id": "c38",
      "from": "n29",
      "to": "n30"
    },
    {
      "id": "c39",
      "from": "n30",
      "to": "n31"
    },
    {
      "id": "c40",
      "from": "n29",
      "to": "n32"
    },
    {
      "id": "c41",
      "from": "n31",
      "to": "n32"
    },
    {
      "id": "c42",
      "from": "n32",
      "to": "n33"
    },
    {
      "id": "c43",
      "from": "n33",
      "to": "n34"
    },
    {
      "id": "c44",
      "from": "n32",
      "to": "n35"
    },
    {
      "id": "c45",
      "from": "n34",
      "to": "n35"
    },
    {
      "id": "c46",
      "from": "n35",
      "to": "n36"
    },
    {
      "id": "c47",
      "from": "n36",
      "to": "n37"
    },
    {
      "id": "c48",
      "from": "n35",
      "to": "n38"
    },
    {
      "id": "c49",
      "from": "n37",
      "to": "n38"
    },
    {
      "id": "c50",
      "from": "n38",
      "to": "n39"
    },
    {
      "id": "c51",
      "from": "n39",
      "to": "n40"
    },
    {
      "id": "c52",
      "from": "n38",
      "to": "n41"
    },
    {
      "id": "c53",
      "from": "n40",
      "to": "n41"
    },
    {
      "id": "c54",
      "from": "n41",
      "to": "n42"
    },
    {
      "id": "c55",
      "from": "n42",
      "to": "n43"
    },
    {
      "id": "c56",
      "from": "n41",
      "to": "n44"
    },
    {
      "id": "c57",
      "from": "n43",
      "to": "n44"
    },
    {
      "id": "c58",
      "from": "n44",
      "to": "n45"
    },
    {
      "id": "c59",
      "from": "n45",
      "to": "n46"
    },
    {
      "id": "c60",
      "from": "n44",
      "to": "n47"
    },
    {
      "id": "c61",
      "from": "n46",
      "to": "n47"
    },
    {
      "id": "c62",
      "from": "n47",
      "to": "n48"
    },
    {
      "id": "c63",
      "from": "n48",
      "to": "n49"
    },
    {
      "id": "c64",
      "from": "n47",
      "to": "n50"
    },
    {
      "id": "c65",
      "from": "n49",
      "to": "n50"
    },
    {
      "id": "c66",
      "from": "n50",
      "to": "n51"
    },
    {
      "id": "c67",
      "from": "n51",
      "to": "n52"
    },
    {
      "id": "c68",
      "from": "n52",
      "to": "n53"
    }
  ]
}