v4: two-lane decoder (two pieces per infer call, grouped-query attention)

README.md

@@ -28,9 +28,11 @@ Neural Engine holds the weights dequantised to fp16, about 4.5 GB, outside the p
 Converted with the pipeline in the VoiceInk fork's `tools/r2t2-coreml`: the encoder through
 coremltools 9 with a fixed 800-frame window (the fused `gelu` replaced by a tanh formulation, whose
 constant absolute error otherwise swamps this encoder's small activations), the decoder through
-[ANEMLL](https://github.com/Anemll/Anemll) as two stateful chunks with a 128-row `prefill` and a
-1-row `infer` function, and a head with a `verify` function that returns the argmax of 128 rows
-in one call. Requires macOS 15 or later (Core ML stateful models) on Apple silicon.
+[ANEMLL](https://github.com/Anemll/Anemll) as two stateful chunks whose KV state holds two
+sequences side by side ("lanes"), with a 128-row `prefill` of one lane and an `infer` that
+decodes one token of each lane in a single pass over the weights, and a head with a two-row
+`infer` and a `verify` function that returns the argmax of 128 rows in one call. Requires macOS 15
+or later (Core ML stateful models) on Apple silicon.
 
 Word error rate, whole utterances, 100-utterance subsets: LibriSpeech test-clean 2.40 %, FLEURS
 pt_br 3.50 %. Against the unconverted bfloat16 weights on MLX, paired on 300 utterances of each:
@@ -41,9 +43,9 @@ pt_br 3.50 %. Against the unconverted bfloat16 weights on MLX, paired on 300 utt
 | file | | size |
 |---|---|---|
 | `R2T2AudioEncoder.mlmodelc` | encoder, fp16, `[1, 128, 800]` mel window + key mask → `[1, 104, 2048]` | 607 MB |
-| `r2t2_FFN_PF_lut8_chunk_01of02.mlmodelc` | decoder layers 0–13, functions `prefill` (128 rows) and `infer` (1 row) | 692 MB |
+| `r2t2_FFN_PF_lut8_chunk_01of02.mlmodelc` | decoder layers 0–13, functions `prefill` (128 rows of one lane) and `infer` (one row per lane) | 692 MB |
 | `r2t2_FFN_PF_lut8_chunk_02of02.mlmodelc` | decoder layers 14–27 and the final norm, same functions | 692 MB |
-| `r2t2_lm_head_lut8.mlmodelc` | head, 16-way split: `infer` (one row → `logits1…16`), `verify` (128 rows → argmaxes) | 306 MB |
+| `r2t2_lm_head_lut8.mlmodelc` | head, 16-way split: `infer` (two rows → `logits1…16`), `verify` (128 rows → argmaxes) | 306 MB |
 | `embed_tokens.f16.bin` | token embeddings, 151 936 × 2048 fp16, row-major, no header | 622 MB |
 | `tokenizer.json`, `tokenizer_config.json` | the original tokenizer | 11 MB |
 | `MODEL_LICENSE`, `LICENSE-Qwen3-ASR.txt`, `NOTICE`, `SHA256SUMS` | | |
@@ -55,17 +57,25 @@ after the licence is accepted. The runtime re-decodes the whole current piece of
 on every pass, so the live text converges on the same result as an offline decode; what makes that
 cheap is keeping the prompt prefix and every completed encoder window in the decoder's KV state,
 prefilling only the new rows plus the previous pass's tokens as a draft, confirming the draft with
-one `verify` call and decoding only from the first divergence. Passes take 150–250 ms on an M5 Pro;
-a 30 s piece decodes offline in about 2.5 s.
+one `verify` call and decoding only from the first divergence. Passes take 150–250 ms on an M5 Pro.
+Offline, two 30 s pieces decode together, one token of each per `infer` call: a 120 s recording
+takes about 6 s (8.7 ms per token; 17 ms for a single piece).
 
 For another runtime, the interface:
 
-- Decoder inputs `hidden_states [1, B, 2048]` fp16, `position_ids [B]` int32,
-  `causal_mask [1, 1, B, 1024]` fp16 (0 to attend, −10 000 otherwise), `current_pos [1]` int32;
-  output `output_hidden_states [1, B, 2048]`. Both chunks share one `MLState` of shape
-  `(56, 8, 1024, 128)` fp16, written by absolute position, so a caller can rewind and overwrite.
-- Head `verify` output per row and slice: `argmax_val`, `argmax_hi`, `argmax_lo` (`[128, 16]` fp16),
-  the slice's best logit and its index as `hi × 64 + lo`.
+- Both chunks share one `MLState` of shape `(56, 8, 2048, 128)` fp16: lane *l* owns cache positions
+  `l·1024 ..< (l+1)·1024`, written by absolute position, so a caller can rewind and overwrite. The
+  causal mask is 0 to attend, −10 000 otherwise; output `output_hidden_states [1, B, 2048]`.
+  - `prefill`: `hidden_states [1, 128, 2048]` fp16, `position_ids [128]` int32 (the lane's own
+    positions), `causal_mask [1, 1, 128, 1024]`, `current_pos [1]` int32 (absolute: `l·1024 + p`),
+    `lane_weights [2, 1, 1]` fp16 one-hot selecting the lane.
+  - `infer`: `hidden_states [1, 2, 2048]` (one row per lane), `position_ids [2]`,
+    `causal_mask [1, 1, 4, 1024]` (each lane's row twice: the two query heads of a K/V head share a
+    matmul row axis), `current_pos [1]` and `current_pos_b [1]` (lane 0 and lane 1, absolute). An
+    idle lane takes a zero row, a fully masked row and a scratch position its next real row overwrites.
+- Head `infer` takes `hidden_states [1, 2, 2048]` and returns `logits1…16` for both rows; `verify`
+  takes 128 rows and returns per row and slice `argmax_val`, `argmax_hi`, `argmax_lo` (`[128, 16]`
+  fp16), the slice's best logit and its index as `hi × 64 + lo`.
 - Prompt: the Qwen3-ASR chat template, with the encoder rows in place of the `<|audio_pad|>`
   embeddings and, to force a language, the assistant turn opened with `language <Name><asr_text>`.
   The model's `|` marks where it stops trusting its own output.

SHA256SUMS

@@ -1,7 +1,7 @@
 0499602714160467f2d68b910651d6216020689f1e016be87a2d0019ee3baeab  tokenizer.json
 4942d005604266809309cabc9f4e9cb89ce855d59b14681fdc0e1cc62ea26c4c  tokenizer_config.json
 efa87abb745fd48d9ee657d62d1a18899e4ee8eba75e47507dbb5991f3b7095a  embed_tokens.f16.bin
-aba6d66c11c29d0da579a44394b9621aca6466611be91053645be6a2d23bbdf4  r2t2_FFN_PF_lut8_chunk_01of02.mlmodelc (tree)
-e30ca12f6ed6b5a0a25af97dfbbedc8c58dfd9db2f707ca130f7941bb12a57eb  r2t2_FFN_PF_lut8_chunk_02of02.mlmodelc (tree)
-ed219b49d16d1bd7905452fdf653eba8778c26d0dd2766d7fc30876eac52cfe5  r2t2_lm_head_lut8.mlmodelc (tree)
+a419770619dc00aaa8d79eb3c4835f162cc2ad275433ef7c56e6663b0dd60da1  r2t2_FFN_PF_lut8_chunk_01of02.mlmodelc (tree)
+9cdb6253b241f4caa3c61c720976ddbedcd92aa0dc9bb933d52c784dd392c5e1  r2t2_FFN_PF_lut8_chunk_02of02.mlmodelc (tree)
+5967ef8ee01652f5b11923afc806774a112cb2f20974d574317138615e0eab9a  r2t2_lm_head_lut8.mlmodelc (tree)
 b34645f8671aec799b836fb4f32ff591d57d061c12ea2c24a3faf0a713b10695  R2T2AudioEncoder.mlmodelc (tree)

r2t2_FFN_PF_lut8_chunk_01of02.mlmodelc/analytics/coremldata.bin

@@ -1,3 +1,3 @@
 version https://git-lfs.github.com/spec/v1
-oid sha256:2cd2e5d8a5bf6e7698165f027d5de6e700c411014d53a8d1db8df6c2fc95e83b
+oid sha256:7174a763588cb29c7e46eacfaa4b0ca20ab99f51f71a53a5e00d8cccbc405732
 size 243

r2t2_FFN_PF_lut8_chunk_01of02.mlmodelc/coremldata.bin

@@ -1,3 +1,3 @@
 version https://git-lfs.github.com/spec/v1
-oid sha256:5b1048f219d90e56c24832df9d307b0c1add26ae76864e5b8cf4075c15ba9a4c
-size 986
+oid sha256:46831adc5fec6999b0fb780725267434239828517c77f9669e94795534afd395
+size 1043

r2t2_FFN_PF_lut8_chunk_01of02.mlmodelc/metadata.json

@@ -2,15 +2,15 @@
   {
     "metadataOutputVersion" : "3.0",
     "userDefinedMetadata" : {
+      "com.github.apple.coremltools.source" : "torch==2.8.0",
+      "com.anemll.lut_bits" : "8",
       "com.github.apple.coremltools.version" : "9.0",
-      "com.anemll.chunk_no" : "1",
-      "com.github.apple.coremltools.source_dialect" : "TorchScript",
       "com.anemll.context_length" : "1024",
-      "com.github.apple.coremltools.source" : "torch==2.8.0",
+      "com.github.apple.coremltools.source_dialect" : "TorchScript",
       "com.anemll.num_chunks" : "2",
       "com.anemll.batch_size" : "128",
       "com.anemll.info" : "Converted with Anemll v0.1.1",
-      "com.anemll.lut_bits" : "8"
+      "com.anemll.chunk_no" : "1"
     },
     "availability" : {
       "macOS" : "15.0",
@@ -25,9 +25,9 @@
         "hasShapeFlexibility" : "0",
         "isOptional" : "0",
         "dataType" : "Float16",
-        "formattedType" : "MultiArray (Float16 1 × 1 × 2048)",
+        "formattedType" : "MultiArray (Float16 1 × 2 × 2048)",
         "shortDescription" : "",
-        "shape" : "[1, 1, 2048]",
+        "shape" : "[1, 2, 2048]",
         "name" : "hidden_states",
         "type" : "MultiArray"
       },
@@ -35,9 +35,9 @@
         "hasShapeFlexibility" : "0",
         "isOptional" : "0",
         "dataType" : "Int32",
-        "formattedType" : "MultiArray (Int32 1)",
+        "formattedType" : "MultiArray (Int32 2)",
         "shortDescription" : "",
-        "shape" : "[1]",
+        "shape" : "[2]",
         "name" : "position_ids",
         "type" : "MultiArray"
       },
@@ -45,9 +45,9 @@
         "hasShapeFlexibility" : "0",
         "isOptional" : "0",
         "dataType" : "Float16",
-        "formattedType" : "MultiArray (Float16 1 × 1 × 1 × 1024)",
+        "formattedType" : "MultiArray (Float16 1 × 1 × 4 × 1024)",
         "shortDescription" : "",
-        "shape" : "[1, 1, 1, 1024]",
+        "shape" : "[1, 1, 4, 1024]",
         "name" : "causal_mask",
         "type" : "MultiArray"
       },
@@ -60,6 +60,16 @@
         "shape" : "[1]",
         "name" : "current_pos",
         "type" : "MultiArray"
+      },
+      {
+        "hasShapeFlexibility" : "0",
+        "isOptional" : "0",
+        "dataType" : "Int32",
+        "formattedType" : "MultiArray (Int32 1)",
+        "shortDescription" : "",
+        "shape" : "[1]",
+        "name" : "current_pos_b",
+        "type" : "MultiArray"
       }
     ],
     "outputSchema" : [
@@ -67,9 +77,9 @@
         "hasShapeFlexibility" : "0",
         "isOptional" : "0",
         "dataType" : "Float16",
-        "formattedType" : "MultiArray (Float16 1 × 1 × 2048)",
+        "formattedType" : "MultiArray (Float16 1 × 2 × 2048)",
         "shortDescription" : "",
-        "shape" : "[1, 1, 2048]",
+        "shape" : "[1, 2, 2048]",
         "name" : "output_hidden_states",
         "type" : "MultiArray"
       }
@@ -86,9 +96,9 @@
             "hasShapeFlexibility" : "0",
             "isOptional" : "0",
             "dataType" : "Float16",
-            "formattedType" : "MultiArray (Float16 1 × 1 × 2048)",
+            "formattedType" : "MultiArray (Float16 1 × 2 × 2048)",
             "shortDescription" : "",
-            "shape" : "[1, 1, 2048]",
+            "shape" : "[1, 2, 2048]",
             "name" : "hidden_states",
             "type" : "MultiArray"
           },
@@ -96,9 +106,9 @@
             "hasShapeFlexibility" : "0",
             "isOptional" : "0",
             "dataType" : "Int32",
-            "formattedType" : "MultiArray (Int32 1)",
+            "formattedType" : "MultiArray (Int32 2)",
             "shortDescription" : "",
-            "shape" : "[1]",
+            "shape" : "[2]",
             "name" : "position_ids",
             "type" : "MultiArray"
           },
@@ -106,9 +116,9 @@
             "hasShapeFlexibility" : "0",
             "isOptional" : "0",
             "dataType" : "Float16",
-            "formattedType" : "MultiArray (Float16 1 × 1 × 1 × 1024)",
+            "formattedType" : "MultiArray (Float16 1 × 1 × 4 × 1024)",
             "shortDescription" : "",
-            "shape" : "[1, 1, 1, 1024]",
+            "shape" : "[1, 1, 4, 1024]",
             "name" : "causal_mask",
             "type" : "MultiArray"
           },
@@ -121,6 +131,16 @@
             "shape" : "[1]",
             "name" : "current_pos",
             "type" : "MultiArray"
+          },
+          {
+            "hasShapeFlexibility" : "0",
+            "isOptional" : "0",
+            "dataType" : "Int32",
+            "formattedType" : "MultiArray (Int32 1)",
+            "shortDescription" : "",
+            "shape" : "[1]",
+            "name" : "current_pos_b",
+            "type" : "MultiArray"
           }
         ],
         "computePrecision" : "Mixed (Float16, Int16, Int32, UInt16)",
@@ -129,9 +149,9 @@
           {
             "dataType" : "Float16",
             "isOptional" : "0",
-            "formattedType" : "State (Float16 56 × 8 × 1024 × 128)",
+            "formattedType" : "State (Float16 56 × 8 × 2048 × 128)",
             "shortDescription" : "",
-            "shape" : "[56, 8, 1024, 128]",
+            "shape" : "[56, 8, 2048, 128]",
             "name" : "model_model_kv_cache_0",
             "type" : "State"
           }
@@ -141,38 +161,36 @@
             "hasShapeFlexibility" : "0",
             "isOptional" : "0",
             "dataType" : "Float16",
-            "formattedType" : "MultiArray (Float16 1 × 1 × 2048)",
+            "formattedType" : "MultiArray (Float16 1 × 2 × 2048)",
             "shortDescription" : "",
-            "shape" : "[1, 1, 2048]",
+            "shape" : "[1, 2, 2048]",
             "name" : "output_hidden_states",
             "type" : "MultiArray"
           }
         ],
         "name" : "infer",
         "mlProgramOperationTypeHistogram" : {
-          "Ios18.expandDims" : 56,
-          "Ios18.mul" : 224,
-          "Ios18.softmax" : 14,
-          "Ios18.matmul" : 28,
-          "Identity" : 1,
+          "Ios18.expandDims" : 84,
+          "Ios18.softmax" : 28,
+          "Ios18.mul" : 238,
+          "Ios18.matmul" : 56,
           "Ios18.greaterEqual" : 2,
           "Select" : 2,
-          "Ios18.readState" : 29,
-          "Tile" : 28,
+          "Ios18.readState" : 57,
           "Ios18.gather" : 2,
-          "Ios18.add" : 73,
+          "Ios18.add" : 88,
           "Ios18.layerNorm" : 56,
-          "Ios18.sliceUpdate" : 28,
-          "Ios18.writeState" : 28,
-          "Ios18.reshape" : 86,
+          "Ios18.sliceUpdate" : 56,
+          "Ios18.writeState" : 56,
+          "Ios18.reshape" : 100,
           "Ios18.constexprLutToDense" : 98,
           "Ios18.conv" : 98,
-          "Ios18.concat" : 140,
-          "Ios18.transpose" : 84,
-          "Ios18.cast" : 5,
+          "Ios18.concat" : 210,
           "Ios18.silu" : 14,
-          "Ios18.sliceByIndex" : 140,
-          "Ios18.squeeze" : 42
+          "Ios18.transpose" : 114,
+          "Ios18.cast" : 5,
+          "Ios18.sliceByIndex" : 282,
+          "Ios18.squeeze" : 14
         }
       },
       {
@@ -216,6 +234,16 @@
             "shape" : "[1]",
             "name" : "current_pos",
             "type" : "MultiArray"
+          },
+          {
+            "hasShapeFlexibility" : "0",
+            "isOptional" : "0",
+            "dataType" : "Float16",
+            "formattedType" : "MultiArray (Float16 2 × 1 × 1)",
+            "shortDescription" : "",
+            "shape" : "[2, 1, 1]",
+            "name" : "lane_weights",
+            "type" : "MultiArray"
           }
         ],
         "computePrecision" : "Mixed (Float16, Int16, Int32, UInt16)",
@@ -224,9 +252,9 @@
           {
             "dataType" : "Float16",
             "isOptional" : "0",
-            "formattedType" : "State (Float16 56 × 8 × 1024 × 128)",
+            "formattedType" : "State (Float16 56 × 8 × 2048 × 128)",
             "shortDescription" : "",
-            "shape" : "[56, 8, 1024, 128]",
+            "shape" : "[56, 8, 2048, 128]",
             "name" : "model_model_kv_cache_0",
             "type" : "State"
           }
@@ -245,27 +273,27 @@
         ],
         "name" : "prefill",
         "mlProgramOperationTypeHistogram" : {
-          "Ios18.expandDims" : 56,
-          "Ios18.mul" : 224,
+          "Ios18.expandDims" : 57,
+          "Ios18.mul" : 280,
           "Ios18.softmax" : 14,
           "Ios18.matmul" : 28,
           "Ios18.greaterEqual" : 2,
           "Select" : 2,
           "Ios18.readState" : 29,
-          "Tile" : 28,
+          "Tile" : 1,
           "Ios18.gather" : 2,
-          "Ios18.add" : 73,
+          "Ios18.add" : 101,
           "Ios18.layerNorm" : 56,
           "Ios18.sliceUpdate" : 28,
           "Ios18.writeState" : 28,
-          "Ios18.reshape" : 114,
+          "Ios18.reshape" : 87,
           "Ios18.constexprLutToDense" : 98,
           "Ios18.conv" : 98,
           "Ios18.concat" : 140,
           "Ios18.transpose" : 128,
           "Ios18.cast" : 5,
           "Ios18.silu" : 14,
-          "Ios18.sliceByIndex" : 140,
+          "Ios18.sliceByIndex" : 198,
           "Ios18.squeeze" : 42
         }
       }
@@ -278,38 +306,36 @@
       {
         "dataType" : "Float16",
         "isOptional" : "0",
-        "formattedType" : "State (Float16 56 × 8 × 1024 × 128)",
+        "formattedType" : "State (Float16 56 × 8 × 2048 × 128)",
         "shortDescription" : "",
-        "shape" : "[56, 8, 1024, 128]",
+        "shape" : "[56, 8, 2048, 128]",
         "name" : "model_model_kv_cache_0",
         "type" : "State"
       }
     ],
     "computePrecision" : "Mixed (Float16, Int16, Int32, UInt16)",
     "mlProgramOperationTypeHistogram" : {
-      "Ios18.expandDims" : 56,
-      "Ios18.mul" : 224,
-      "Ios18.softmax" : 14,
-      "Ios18.matmul" : 28,
-      "Identity" : 1,
+      "Ios18.expandDims" : 84,
+      "Ios18.softmax" : 28,
+      "Ios18.mul" : 238,
+      "Ios18.matmul" : 56,
       "Ios18.greaterEqual" : 2,
       "Select" : 2,
-      "Ios18.readState" : 29,
-      "Tile" : 28,
+      "Ios18.readState" : 57,
       "Ios18.gather" : 2,
-      "Ios18.add" : 73,
+      "Ios18.add" : 88,
       "Ios18.layerNorm" : 56,
-      "Ios18.sliceUpdate" : 28,
-      "Ios18.writeState" : 28,
-      "Ios18.reshape" : 86,
+      "Ios18.sliceUpdate" : 56,
+      "Ios18.writeState" : 56,
+      "Ios18.reshape" : 100,
       "Ios18.constexprLutToDense" : 98,
       "Ios18.conv" : 98,
-      "Ios18.concat" : 140,
-      "Ios18.transpose" : 84,
-      "Ios18.cast" : 5,
+      "Ios18.concat" : 210,
       "Ios18.silu" : 14,
-      "Ios18.sliceByIndex" : 140,
-      "Ios18.squeeze" : 42
+      "Ios18.transpose" : 114,
+      "Ios18.cast" : 5,
+      "Ios18.sliceByIndex" : 282,
+      "Ios18.squeeze" : 14
     },
     "shortDescription" : "Anemll Model: Multifunction FFN+Prefill",
     "generatedClassName" : "r2t2_FFN_PF_lut8_chunk_01of02",

r2t2_FFN_PF_lut8_chunk_01of02.mlmodelc/weights/weight.bin

@@ -1,3 +1,3 @@
 version https://git-lfs.github.com/spec/v1
-oid sha256:74494c85747845fc4f27bfc1c77e6b81ec3ebd0a494aca59c35d59ab7df62485
-size 724179904
+oid sha256:a0f0caf69eaf2aa6e51bcd99a1772c0772ba732dbad4e1aaa0b3d6ed0cd6007a
+size 725228480

r2t2_FFN_PF_lut8_chunk_02of02.mlmodelc/analytics/coremldata.bin

@@ -1,3 +1,3 @@
 version https://git-lfs.github.com/spec/v1
-oid sha256:c90c7e7bf41e01fd443833700072119b021ea2c0816241bc6b303888181931d2
+oid sha256:182f3ffa7c5ab2c90563a7180728f7b5604759464587d081c3ce23326f57d06e
 size 243

r2t2_FFN_PF_lut8_chunk_02of02.mlmodelc/coremldata.bin

@@ -1,3 +1,3 @@
 version https://git-lfs.github.com/spec/v1
-oid sha256:e45f2ea058ea90c17e7ab6d48ab0451cfdc11a1167aa3971ddf32563f742a2d8
-size 986
+oid sha256:085250e41cdcdba9ea12f0ae908a053f8e967d7e244b94f753006c5ff5062f91
+size 1043

r2t2_FFN_PF_lut8_chunk_02of02.mlmodelc/metadata.json

@@ -2,14 +2,14 @@
   {
     "metadataOutputVersion" : "3.0",
     "userDefinedMetadata" : {
+      "com.github.apple.coremltools.source" : "torch==2.8.0",
       "com.github.apple.coremltools.version" : "9.0",
-      "com.github.apple.coremltools.source_dialect" : "TorchScript",
       "com.anemll.context_length" : "1024",
-      "com.github.apple.coremltools.source" : "torch==2.8.0",
+      "com.github.apple.coremltools.source_dialect" : "TorchScript",
       "com.anemll.lut_bits" : "8",
       "com.anemll.num_chunks" : "2",
-      "com.anemll.batch_size" : "128",
       "com.anemll.info" : "Converted with Anemll v0.1.1",
+      "com.anemll.batch_size" : "128",
       "com.anemll.chunk_no" : "2"
     },
     "availability" : {
@@ -25,9 +25,9 @@
         "hasShapeFlexibility" : "0",
         "isOptional" : "0",
         "dataType" : "Float16",
-        "formattedType" : "MultiArray (Float16 1 × 1 × 2048)",
+        "formattedType" : "MultiArray (Float16 1 × 2 × 2048)",
         "shortDescription" : "",
-        "shape" : "[1, 1, 2048]",
+        "shape" : "[1, 2, 2048]",
         "name" : "hidden_states",
         "type" : "MultiArray"
       },
@@ -35,9 +35,9 @@
         "hasShapeFlexibility" : "0",
         "isOptional" : "0",
         "dataType" : "Int32",
-        "formattedType" : "MultiArray (Int32 1)",
+        "formattedType" : "MultiArray (Int32 2)",
         "shortDescription" : "",
-        "shape" : "[1]",
+        "shape" : "[2]",
         "name" : "position_ids",
         "type" : "MultiArray"
       },
@@ -45,9 +45,9 @@
         "hasShapeFlexibility" : "0",
         "isOptional" : "0",
         "dataType" : "Float16",
-        "formattedType" : "MultiArray (Float16 1 × 1 × 1 × 1024)",
+        "formattedType" : "MultiArray (Float16 1 × 1 × 4 × 1024)",
         "shortDescription" : "",
-        "shape" : "[1, 1, 1, 1024]",
+        "shape" : "[1, 1, 4, 1024]",
         "name" : "causal_mask",
         "type" : "MultiArray"
       },
@@ -60,6 +60,16 @@
         "shape" : "[1]",
         "name" : "current_pos",
         "type" : "MultiArray"
+      },
+      {
+        "hasShapeFlexibility" : "0",
+        "isOptional" : "0",
+        "dataType" : "Int32",
+        "formattedType" : "MultiArray (Int32 1)",
+        "shortDescription" : "",
+        "shape" : "[1]",
+        "name" : "current_pos_b",
+        "type" : "MultiArray"
       }
     ],
     "outputSchema" : [
@@ -67,9 +77,9 @@
         "hasShapeFlexibility" : "0",
         "isOptional" : "0",
         "dataType" : "Float16",
-        "formattedType" : "MultiArray (Float16 1 × 1 × 2048)",
+        "formattedType" : "MultiArray (Float16 1 × 2 × 2048)",
         "shortDescription" : "",
-        "shape" : "[1, 1, 2048]",
+        "shape" : "[1, 2, 2048]",
         "name" : "output_hidden_states",
         "type" : "MultiArray"
       }
@@ -86,9 +96,9 @@
             "hasShapeFlexibility" : "0",
             "isOptional" : "0",
             "dataType" : "Float16",
-            "formattedType" : "MultiArray (Float16 1 × 1 × 2048)",
+            "formattedType" : "MultiArray (Float16 1 × 2 × 2048)",
             "shortDescription" : "",
-            "shape" : "[1, 1, 2048]",
+            "shape" : "[1, 2, 2048]",
             "name" : "hidden_states",
             "type" : "MultiArray"
           },
@@ -96,9 +106,9 @@
             "hasShapeFlexibility" : "0",
             "isOptional" : "0",
             "dataType" : "Int32",
-            "formattedType" : "MultiArray (Int32 1)",
+            "formattedType" : "MultiArray (Int32 2)",
             "shortDescription" : "",
-            "shape" : "[1]",
+            "shape" : "[2]",
             "name" : "position_ids",
             "type" : "MultiArray"
           },
@@ -106,9 +116,9 @@
             "hasShapeFlexibility" : "0",
             "isOptional" : "0",
             "dataType" : "Float16",
-            "formattedType" : "MultiArray (Float16 1 × 1 × 1 × 1024)",
+            "formattedType" : "MultiArray (Float16 1 × 1 × 4 × 1024)",
             "shortDescription" : "",
-            "shape" : "[1, 1, 1, 1024]",
+            "shape" : "[1, 1, 4, 1024]",
             "name" : "causal_mask",
             "type" : "MultiArray"
           },
@@ -121,6 +131,16 @@
             "shape" : "[1]",
             "name" : "current_pos",
             "type" : "MultiArray"
+          },
+          {
+            "hasShapeFlexibility" : "0",
+            "isOptional" : "0",
+            "dataType" : "Int32",
+            "formattedType" : "MultiArray (Int32 1)",
+            "shortDescription" : "",
+            "shape" : "[1]",
+            "name" : "current_pos_b",
+            "type" : "MultiArray"
           }
         ],
         "computePrecision" : "Mixed (Float16, Int16, Int32, UInt16)",
@@ -129,9 +149,9 @@
           {
             "dataType" : "Float16",
             "isOptional" : "0",
-            "formattedType" : "State (Float16 56 × 8 × 1024 × 128)",
+            "formattedType" : "State (Float16 56 × 8 × 2048 × 128)",
             "shortDescription" : "",
-            "shape" : "[56, 8, 1024, 128]",
+            "shape" : "[56, 8, 2048, 128]",
             "name" : "model_model_kv_cache_0",
             "type" : "State"
           }
@@ -141,38 +161,36 @@
             "hasShapeFlexibility" : "0",
             "isOptional" : "0",
             "dataType" : "Float16",
-            "formattedType" : "MultiArray (Float16 1 × 1 × 2048)",
+            "formattedType" : "MultiArray (Float16 1 × 2 × 2048)",
             "shortDescription" : "",
-            "shape" : "[1, 1, 2048]",
+            "shape" : "[1, 2, 2048]",
             "name" : "output_hidden_states",
             "type" : "MultiArray"
           }
         ],
         "name" : "infer",
         "mlProgramOperationTypeHistogram" : {
-          "Ios18.expandDims" : 56,
-          "Ios18.mul" : 226,
-          "Ios18.softmax" : 14,
-          "Ios18.matmul" : 28,
-          "Identity" : 1,
+          "Ios18.expandDims" : 84,
+          "Ios18.softmax" : 28,
+          "Ios18.mul" : 240,
+          "Ios18.matmul" : 56,
           "Ios18.greaterEqual" : 2,
           "Select" : 2,
-          "Ios18.readState" : 29,
-          "Tile" : 28,
+          "Ios18.readState" : 57,
           "Ios18.gather" : 2,
-          "Ios18.add" : 73,
+          "Ios18.add" : 88,
           "Ios18.layerNorm" : 57,
-          "Ios18.sliceUpdate" : 28,
-          "Ios18.writeState" : 28,
-          "Ios18.reshape" : 86,
+          "Ios18.sliceUpdate" : 56,
+          "Ios18.writeState" : 56,
+          "Ios18.reshape" : 100,
           "Ios18.constexprLutToDense" : 98,
           "Ios18.conv" : 98,
-          "Ios18.concat" : 141,
-          "Ios18.transpose" : 84,
-          "Ios18.cast" : 5,
+          "Ios18.concat" : 211,
           "Ios18.silu" : 14,
-          "Ios18.sliceByIndex" : 141,
-          "Ios18.squeeze" : 42
+          "Ios18.transpose" : 114,
+          "Ios18.cast" : 5,
+          "Ios18.sliceByIndex" : 283,
+          "Ios18.squeeze" : 14
         }
       },
       {
@@ -216,6 +234,16 @@
             "shape" : "[1]",
             "name" : "current_pos",
             "type" : "MultiArray"
+          },
+          {
+            "hasShapeFlexibility" : "0",
+            "isOptional" : "0",
+            "dataType" : "Float16",
+            "formattedType" : "MultiArray (Float16 2 × 1 × 1)",
+            "shortDescription" : "",
+            "shape" : "[2, 1, 1]",
+            "name" : "lane_weights",
+            "type" : "MultiArray"
           }
         ],
         "computePrecision" : "Mixed (Float16, Int16, Int32, UInt16)",
@@ -224,9 +252,9 @@
           {
             "dataType" : "Float16",
             "isOptional" : "0",
-            "formattedType" : "State (Float16 56 × 8 × 1024 × 128)",
+            "formattedType" : "State (Float16 56 × 8 × 2048 × 128)",
             "shortDescription" : "",
-            "shape" : "[56, 8, 1024, 128]",
+            "shape" : "[56, 8, 2048, 128]",
             "name" : "model_model_kv_cache_0",
             "type" : "State"
           }
@@ -245,27 +273,27 @@
         ],
         "name" : "prefill",
         "mlProgramOperationTypeHistogram" : {
-          "Ios18.expandDims" : 56,
-          "Ios18.mul" : 226,
+          "Ios18.expandDims" : 57,
+          "Ios18.mul" : 282,
           "Ios18.softmax" : 14,
           "Ios18.matmul" : 28,
           "Ios18.greaterEqual" : 2,
           "Select" : 2,
           "Ios18.readState" : 29,
-          "Tile" : 28,
+          "Tile" : 1,
           "Ios18.gather" : 2,
-          "Ios18.add" : 73,
+          "Ios18.add" : 101,
           "Ios18.layerNorm" : 57,
           "Ios18.sliceUpdate" : 28,
           "Ios18.writeState" : 28,
-          "Ios18.reshape" : 114,
+          "Ios18.reshape" : 87,
           "Ios18.constexprLutToDense" : 98,
           "Ios18.conv" : 98,
           "Ios18.concat" : 141,
           "Ios18.transpose" : 128,
           "Ios18.cast" : 5,
           "Ios18.silu" : 14,
-          "Ios18.sliceByIndex" : 141,
+          "Ios18.sliceByIndex" : 199,
           "Ios18.squeeze" : 42
         }
       }
@@ -278,38 +306,36 @@
       {
         "dataType" : "Float16",
         "isOptional" : "0",
-        "formattedType" : "State (Float16 56 × 8 × 1024 × 128)",
+        "formattedType" : "State (Float16 56 × 8 × 2048 × 128)",
         "shortDescription" : "",
-        "shape" : "[56, 8, 1024, 128]",
+        "shape" : "[56, 8, 2048, 128]",
         "name" : "model_model_kv_cache_0",
         "type" : "State"
       }
     ],
     "computePrecision" : "Mixed (Float16, Int16, Int32, UInt16)",
     "mlProgramOperationTypeHistogram" : {
-      "Ios18.expandDims" : 56,
-      "Ios18.mul" : 226,
-      "Ios18.softmax" : 14,
-      "Ios18.matmul" : 28,
-      "Identity" : 1,
+      "Ios18.expandDims" : 84,
+      "Ios18.softmax" : 28,
+      "Ios18.mul" : 240,
+      "Ios18.matmul" : 56,
       "Ios18.greaterEqual" : 2,
       "Select" : 2,
-      "Ios18.readState" : 29,
-      "Tile" : 28,
+      "Ios18.readState" : 57,
       "Ios18.gather" : 2,
-      "Ios18.add" : 73,
+      "Ios18.add" : 88,
       "Ios18.layerNorm" : 57,
-      "Ios18.sliceUpdate" : 28,
-      "Ios18.writeState" : 28,
-      "Ios18.reshape" : 86,
+      "Ios18.sliceUpdate" : 56,
+      "Ios18.writeState" : 56,
+      "Ios18.reshape" : 100,
       "Ios18.constexprLutToDense" : 98,
       "Ios18.conv" : 98,
-      "Ios18.concat" : 141,
-      "Ios18.transpose" : 84,
-      "Ios18.cast" : 5,
+      "Ios18.concat" : 211,
       "Ios18.silu" : 14,
-      "Ios18.sliceByIndex" : 141,
-      "Ios18.squeeze" : 42
+      "Ios18.transpose" : 114,
+      "Ios18.cast" : 5,
+      "Ios18.sliceByIndex" : 283,
+      "Ios18.squeeze" : 14
     },
     "shortDescription" : "Anemll Model: Multifunction FFN+Prefill",
     "generatedClassName" : "r2t2_FFN_PF_lut8_chunk_02of02",

r2t2_FFN_PF_lut8_chunk_02of02.mlmodelc/weights/weight.bin

@@ -1,3 +1,3 @@
 version https://git-lfs.github.com/spec/v1
-oid sha256:d2132cc8f5338b94f1b568c3df1d821d8b99db829cab291c56d56f6baaa0b20b
-size 724184064
+oid sha256:37c6f521391020156c1ff2a1a652e8696cdc555acdc1c8e56db2f5569cbf2c03
+size 725232640

r2t2_lm_head_lut8.mlmodelc/analytics/coremldata.bin

@@ -1,3 +1,3 @@
 version https://git-lfs.github.com/spec/v1
-oid sha256:3c98575b5a199847eb87ba56275e4b3f58bf6d3b28d9275b6b51334e27a7efbb
+oid sha256:1273e52fd5323604f856809221c81bd50bb513460b1d70838ef0ca18a598ef93
 size 243

r2t2_lm_head_lut8.mlmodelc/coremldata.bin

@@ -1,3 +1,3 @@
 version https://git-lfs.github.com/spec/v1
-oid sha256:552fdb564ca00497c313551966843e18eb04616c0a24190e64e0db1c5358a8da
-size 1264
+oid sha256:5d35d6cfd446c0980ef31fc5ec454219e5316dffb4f0fb89f7cd8191ead3cc5d
+size 1262

r2t2_lm_head_lut8.mlmodelc/metadata.json

@@ -3,12 +3,12 @@
     "metadataOutputVersion" : "3.0",
     "userDefinedMetadata" : {
       "com.github.apple.coremltools.version" : "9.0",
+      "com.github.apple.coremltools.source_dialect" : "TorchScript",
       "com.github.apple.coremltools.source" : "torch==2.8.0",
       "com.anemll.context_length" : "1024",
-      "com.github.apple.coremltools.source_dialect" : "TorchScript",
       "com.anemll.lm_head_chunk_sizes" : "9496,9496,9496,9496,9496,9496,9496,9496,9496,9496,9496,9496,9496,9496,9496,9496",
       "com.anemll.vocab_size" : "151936",
-      "com.anemll.batch_size" : "128",
+      "com.anemll.batch_size" : "2",
       "com.anemll.info" : "Converted with Anemll v0.1.1",
       "com.anemll.argmax_in_model" : "true",
       "com.anemll.lut_bits" : "8"
@@ -26,9 +26,9 @@
         "hasShapeFlexibility" : "0",
         "isOptional" : "0",
         "dataType" : "Float16",
-        "formattedType" : "MultiArray (Float16 1 × 1 × 2048)",
+        "formattedType" : "MultiArray (Float16 1 × 2 × 2048)",
         "shortDescription" : "",
-        "shape" : "[1, 1, 2048]",
+        "shape" : "[1, 2, 2048]",
         "name" : "hidden_states",
         "type" : "MultiArray"
       }
@@ -38,9 +38,9 @@
         "hasShapeFlexibility" : "0",
         "isOptional" : "0",
         "dataType" : "Float16",
-        "formattedType" : "MultiArray (Float16 1 × 1 × 9496)",
+        "formattedType" : "MultiArray (Float16 1 × 2 × 9496)",
         "shortDescription" : "",
-        "shape" : "[1, 1, 9496]",
+        "shape" : "[1, 2, 9496]",
         "name" : "logits1",
         "type" : "MultiArray"
       },
@@ -48,9 +48,9 @@
         "hasShapeFlexibility" : "0",
         "isOptional" : "0",
         "dataType" : "Float16",
-        "formattedType" : "MultiArray (Float16 1 × 1 × 9496)",
+        "formattedType" : "MultiArray (Float16 1 × 2 × 9496)",
         "shortDescription" : "",
-        "shape" : "[1, 1, 9496]",
+        "shape" : "[1, 2, 9496]",
         "name" : "logits2",
         "type" : "MultiArray"
       },
@@ -58,9 +58,9 @@
         "hasShapeFlexibility" : "0",
         "isOptional" : "0",
         "dataType" : "Float16",
-        "formattedType" : "MultiArray (Float16 1 × 1 × 9496)",
+        "formattedType" : "MultiArray (Float16 1 × 2 × 9496)",
         "shortDescription" : "",
-        "shape" : "[1, 1, 9496]",
+        "shape" : "[1, 2, 9496]",
         "name" : "logits3",
         "type" : "MultiArray"
       },
@@ -68,9 +68,9 @@
         "hasShapeFlexibility" : "0",
         "isOptional" : "0",
         "dataType" : "Float16",
-        "formattedType" : "MultiArray (Float16 1 × 1 × 9496)",
+        "formattedType" : "MultiArray (Float16 1 × 2 × 9496)",
         "shortDescription" : "",
-        "shape" : "[1, 1, 9496]",
+        "shape" : "[1, 2, 9496]",
         "name" : "logits4",
         "type" : "MultiArray"
       },
@@ -78,9 +78,9 @@
         "hasShapeFlexibility" : "0",
         "isOptional" : "0",
         "dataType" : "Float16",
-        "formattedType" : "MultiArray (Float16 1 × 1 × 9496)",
+        "formattedType" : "MultiArray (Float16 1 × 2 × 9496)",
         "shortDescription" : "",
-        "shape" : "[1, 1, 9496]",
+        "shape" : "[1, 2, 9496]",
         "name" : "logits5",
         "type" : "MultiArray"
       },
@@ -88,9 +88,9 @@
         "hasShapeFlexibility" : "0",
         "isOptional" : "0",
         "dataType" : "Float16",
-        "formattedType" : "MultiArray (Float16 1 × 1 × 9496)",
+        "formattedType" : "MultiArray (Float16 1 × 2 × 9496)",
         "shortDescription" : "",
-        "shape" : "[1, 1, 9496]",
+        "shape" : "[1, 2, 9496]",
         "name" : "logits6",
         "type" : "MultiArray"
       },
@@ -98,9 +98,9 @@
         "hasShapeFlexibility" : "0",
         "isOptional" : "0",
         "dataType" : "Float16",
-        "formattedType" : "MultiArray (Float16 1 × 1 × 9496)",
+        "formattedType" : "MultiArray (Float16 1 × 2 × 9496)",
         "shortDescription" : "",
-        "shape" : "[1, 1, 9496]",
+        "shape" : "[1, 2, 9496]",
         "name" : "logits7",
         "type" : "MultiArray"
       },
@@ -108,9 +108,9 @@
         "hasShapeFlexibility" : "0",
         "isOptional" : "0",
         "dataType" : "Float16",
-        "formattedType" : "MultiArray (Float16 1 × 1 × 9496)",
+        "formattedType" : "MultiArray (Float16 1 × 2 × 9496)",
         "shortDescription" : "",
-        "shape" : "[1, 1, 9496]",
+        "shape" : "[1, 2, 9496]",
         "name" : "logits8",
         "type" : "MultiArray"
       },
@@ -118,9 +118,9 @@
         "hasShapeFlexibility" : "0",
         "isOptional" : "0",
         "dataType" : "Float16",
-        "formattedType" : "MultiArray (Float16 1 × 1 × 9496)",
+        "formattedType" : "MultiArray (Float16 1 × 2 × 9496)",
         "shortDescription" : "",
-        "shape" : "[1, 1, 9496]",
+        "shape" : "[1, 2, 9496]",
         "name" : "logits9",
         "type" : "MultiArray"
       },
@@ -128,9 +128,9 @@
         "hasShapeFlexibility" : "0",
         "isOptional" : "0",
         "dataType" : "Float16",
-        "formattedType" : "MultiArray (Float16 1 × 1 × 9496)",
+        "formattedType" : "MultiArray (Float16 1 × 2 × 9496)",
         "shortDescription" : "",
-        "shape" : "[1, 1, 9496]",
+        "shape" : "[1, 2, 9496]",
         "name" : "logits10",
         "type" : "MultiArray"
       },
@@ -138,9 +138,9 @@
         "hasShapeFlexibility" : "0",
         "isOptional" : "0",
         "dataType" : "Float16",
-        "formattedType" : "MultiArray (Float16 1 × 1 × 9496)",
+        "formattedType" : "MultiArray (Float16 1 × 2 × 9496)",
         "shortDescription" : "",
-        "shape" : "[1, 1, 9496]",
+        "shape" : "[1, 2, 9496]",
         "name" : "logits11",
         "type" : "MultiArray"
       },
@@ -148,9 +148,9 @@
         "hasShapeFlexibility" : "0",
         "isOptional" : "0",
         "dataType" : "Float16",
-        "formattedType" : "MultiArray (Float16 1 × 1 × 9496)",
+        "formattedType" : "MultiArray (Float16 1 × 2 × 9496)",
         "shortDescription" : "",
-        "shape" : "[1, 1, 9496]",
+        "shape" : "[1, 2, 9496]",
         "name" : "logits12",
         "type" : "MultiArray"
       },
@@ -158,9 +158,9 @@
         "hasShapeFlexibility" : "0",
         "isOptional" : "0",
         "dataType" : "Float16",
-        "formattedType" : "MultiArray (Float16 1 × 1 × 9496)",
+        "formattedType" : "MultiArray (Float16 1 × 2 × 9496)",
         "shortDescription" : "",
-        "shape" : "[1, 1, 9496]",
+        "shape" : "[1, 2, 9496]",
         "name" : "logits13",
         "type" : "MultiArray"
       },
@@ -168,9 +168,9 @@
         "hasShapeFlexibility" : "0",
         "isOptional" : "0",
         "dataType" : "Float16",
-        "formattedType" : "MultiArray (Float16 1 × 1 × 9496)",
+        "formattedType" : "MultiArray (Float16 1 × 2 × 9496)",
         "shortDescription" : "",
-        "shape" : "[1, 1, 9496]",
+        "shape" : "[1, 2, 9496]",
         "name" : "logits14",
         "type" : "MultiArray"
       },
@@ -178,9 +178,9 @@
         "hasShapeFlexibility" : "0",
         "isOptional" : "0",
         "dataType" : "Float16",
-        "formattedType" : "MultiArray (Float16 1 × 1 × 9496)",
+        "formattedType" : "MultiArray (Float16 1 × 2 × 9496)",
         "shortDescription" : "",
-        "shape" : "[1, 1, 9496]",
+        "shape" : "[1, 2, 9496]",
         "name" : "logits15",
         "type" : "MultiArray"
       },
@@ -188,9 +188,9 @@
         "hasShapeFlexibility" : "0",
         "isOptional" : "0",
         "dataType" : "Float16",
-        "formattedType" : "MultiArray (Float16 1 × 1 × 9496)",
+        "formattedType" : "MultiArray (Float16 1 × 2 × 9496)",
         "shortDescription" : "",
-        "shape" : "[1, 1, 9496]",
+        "shape" : "[1, 2, 9496]",
         "name" : "logits16",
         "type" : "MultiArray"
       }
@@ -207,9 +207,9 @@
             "hasShapeFlexibility" : "0",
             "isOptional" : "0",
             "dataType" : "Float16",
-            "formattedType" : "MultiArray (Float16 1 × 1 × 2048)",
+            "formattedType" : "MultiArray (Float16 1 × 2 × 2048)",
             "shortDescription" : "",
-            "shape" : "[1, 1, 2048]",
+            "shape" : "[1, 2, 2048]",
             "name" : "hidden_states",
             "type" : "MultiArray"
           }
@@ -224,9 +224,9 @@
             "hasShapeFlexibility" : "0",
             "isOptional" : "0",
             "dataType" : "Float16",
-            "formattedType" : "MultiArray (Float16 1 × 1 × 9496)",
+            "formattedType" : "MultiArray (Float16 1 × 2 × 9496)",
             "shortDescription" : "",
-            "shape" : "[1, 1, 9496]",
+            "shape" : "[1, 2, 9496]",
             "name" : "logits1",
             "type" : "MultiArray"
           },
@@ -234,9 +234,9 @@
             "hasShapeFlexibility" : "0",
             "isOptional" : "0",
             "dataType" : "Float16",
-            "formattedType" : "MultiArray (Float16 1 × 1 × 9496)",
+            "formattedType" : "MultiArray (Float16 1 × 2 × 9496)",
             "shortDescription" : "",
-            "shape" : "[1, 1, 9496]",
+            "shape" : "[1, 2, 9496]",
             "name" : "logits2",
             "type" : "MultiArray"
           },
@@ -244,9 +244,9 @@
             "hasShapeFlexibility" : "0",
             "isOptional" : "0",
             "dataType" : "Float16",
-            "formattedType" : "MultiArray (Float16 1 × 1 × 9496)",
+            "formattedType" : "MultiArray (Float16 1 × 2 × 9496)",
             "shortDescription" : "",
-            "shape" : "[1, 1, 9496]",
+            "shape" : "[1, 2, 9496]",
             "name" : "logits3",
             "type" : "MultiArray"
           },
@@ -254,9 +254,9 @@
             "hasShapeFlexibility" : "0",
             "isOptional" : "0",
             "dataType" : "Float16",
-            "formattedType" : "MultiArray (Float16 1 × 1 × 9496)",
+            "formattedType" : "MultiArray (Float16 1 × 2 × 9496)",
             "shortDescription" : "",
-            "shape" : "[1, 1, 9496]",
+            "shape" : "[1, 2, 9496]",
             "name" : "logits4",
             "type" : "MultiArray"
           },
@@ -264,9 +264,9 @@
             "hasShapeFlexibility" : "0",
             "isOptional" : "0",
             "dataType" : "Float16",
-            "formattedType" : "MultiArray (Float16 1 × 1 × 9496)",
+            "formattedType" : "MultiArray (Float16 1 × 2 × 9496)",
             "shortDescription" : "",
-            "shape" : "[1, 1, 9496]",
+            "shape" : "[1, 2, 9496]",
             "name" : "logits5",
             "type" : "MultiArray"
           },
@@ -274,9 +274,9 @@
             "hasShapeFlexibility" : "0",
             "isOptional" : "0",
             "dataType" : "Float16",
-            "formattedType" : "MultiArray (Float16 1 × 1 × 9496)",
+            "formattedType" : "MultiArray (Float16 1 × 2 × 9496)",
             "shortDescription" : "",
-            "shape" : "[1, 1, 9496]",
+            "shape" : "[1, 2, 9496]",
             "name" : "logits6",
             "type" : "MultiArray"
           },
@@ -284,9 +284,9 @@
             "hasShapeFlexibility" : "0",
             "isOptional" : "0",
             "dataType" : "Float16",
-            "formattedType" : "MultiArray (Float16 1 × 1 × 9496)",
+            "formattedType" : "MultiArray (Float16 1 × 2 × 9496)",
             "shortDescription" : "",
-            "shape" : "[1, 1, 9496]",
+            "shape" : "[1, 2, 9496]",
             "name" : "logits7",
             "type" : "MultiArray"
           },
@@ -294,9 +294,9 @@
             "hasShapeFlexibility" : "0",
             "isOptional" : "0",
             "dataType" : "Float16",
-            "formattedType" : "MultiArray (Float16 1 × 1 × 9496)",
+            "formattedType" : "MultiArray (Float16 1 × 2 × 9496)",
             "shortDescription" : "",
-            "shape" : "[1, 1, 9496]",
+            "shape" : "[1, 2, 9496]",
             "name" : "logits8",
             "type" : "MultiArray"
           },
@@ -304,9 +304,9 @@
             "hasShapeFlexibility" : "0",
             "isOptional" : "0",
             "dataType" : "Float16",
-            "formattedType" : "MultiArray (Float16 1 × 1 × 9496)",
+            "formattedType" : "MultiArray (Float16 1 × 2 × 9496)",
             "shortDescription" : "",
-            "shape" : "[1, 1, 9496]",
+            "shape" : "[1, 2, 9496]",
             "name" : "logits9",
             "type" : "MultiArray"
           },
@@ -314,9 +314,9 @@
             "hasShapeFlexibility" : "0",
             "isOptional" : "0",
             "dataType" : "Float16",
-            "formattedType" : "MultiArray (Float16 1 × 1 × 9496)",
+            "formattedType" : "MultiArray (Float16 1 × 2 × 9496)",
             "shortDescription" : "",
-            "shape" : "[1, 1, 9496]",
+            "shape" : "[1, 2, 9496]",
             "name" : "logits10",
             "type" : "MultiArray"
           },
@@ -324,9 +324,9 @@
             "hasShapeFlexibility" : "0",
             "isOptional" : "0",
             "dataType" : "Float16",
-            "formattedType" : "MultiArray (Float16 1 × 1 × 9496)",
+            "formattedType" : "MultiArray (Float16 1 × 2 × 9496)",
             "shortDescription" : "",
-            "shape" : "[1, 1, 9496]",
+            "shape" : "[1, 2, 9496]",
             "name" : "logits11",
             "type" : "MultiArray"
           },
@@ -334,9 +334,9 @@
             "hasShapeFlexibility" : "0",
             "isOptional" : "0",
             "dataType" : "Float16",
-            "formattedType" : "MultiArray (Float16 1 × 1 × 9496)",
+            "formattedType" : "MultiArray (Float16 1 × 2 × 9496)",
             "shortDescription" : "",
-            "shape" : "[1, 1, 9496]",
+            "shape" : "[1, 2, 9496]",
             "name" : "logits12",
             "type" : "MultiArray"
           },
@@ -344,9 +344,9 @@
             "hasShapeFlexibility" : "0",
             "isOptional" : "0",
             "dataType" : "Float16",
-            "formattedType" : "MultiArray (Float16 1 × 1 × 9496)",
+            "formattedType" : "MultiArray (Float16 1 × 2 × 9496)",
             "shortDescription" : "",
-            "shape" : "[1, 1, 9496]",
+            "shape" : "[1, 2, 9496]",
             "name" : "logits13",
             "type" : "MultiArray"
           },
@@ -354,9 +354,9 @@
             "hasShapeFlexibility" : "0",
             "isOptional" : "0",
             "dataType" : "Float16",
-            "formattedType" : "MultiArray (Float16 1 × 1 × 9496)",
+            "formattedType" : "MultiArray (Float16 1 × 2 × 9496)",
             "shortDescription" : "",
-            "shape" : "[1, 1, 9496]",
+            "shape" : "[1, 2, 9496]",
             "name" : "logits14",
             "type" : "MultiArray"
           },
@@ -364,9 +364,9 @@
             "hasShapeFlexibility" : "0",
             "isOptional" : "0",
             "dataType" : "Float16",
-            "formattedType" : "MultiArray (Float16 1 × 1 × 9496)",
+            "formattedType" : "MultiArray (Float16 1 × 2 × 9496)",
             "shortDescription" : "",
-            "shape" : "[1, 1, 9496]",
+            "shape" : "[1, 2, 9496]",
             "name" : "logits15",
             "type" : "MultiArray"
           },
@@ -374,9 +374,9 @@
             "hasShapeFlexibility" : "0",
             "isOptional" : "0",
             "dataType" : "Float16",
-            "formattedType" : "MultiArray (Float16 1 × 1 × 9496)",
+            "formattedType" : "MultiArray (Float16 1 × 2 × 9496)",
             "shortDescription" : "",
-            "shape" : "[1, 1, 9496]",
+            "shape" : "[1, 2, 9496]",
             "name" : "logits16",
             "type" : "MultiArray"
           }