Faster argmax by skipping logit concat

Llama-3.2-1B-Instruct_chunk6.mlmodelc/analytics/coremldata.bin

@@ -1,3 +1,3 @@
 version https://git-lfs.github.com/spec/v1
-oid sha256:23aa1a8a2b6bca88beeecf08d6281cdeb43aad33eb02cbe10ebff7eede7ed329
+oid sha256:13c3ef8de95508f3e6076553fb737642a123494d4403aef5ebda7ce0c9acc236
 size 243

Llama-3.2-1B-Instruct_chunk6.mlmodelc/coremldata.bin

@@ -1,3 +1,3 @@
 version https://git-lfs.github.com/spec/v1
-oid sha256:2740a88eec733758ae31768866e8550009ce36eb5177a09938d42cb18a095d05
-size 311
+oid sha256:c0f38714827a836f9d17e85d57fd6906b69f0bab1eefba12c1ade1e3eaf9715b
+size 501

Llama-3.2-1B-Instruct_chunk6.mlmodelc/metadata.json

@@ -7,10 +7,80 @@
         "hasShapeFlexibility" : "0",
         "isOptional" : "0",
         "dataType" : "Float16",
-        "formattedType" : "MultiArray (Float16 1 × 64 × 128256)",
+        "formattedType" : "MultiArray (Float16 1 × 64 × 16384)",
         "shortDescription" : "",
-        "shape" : "[1, 64, 128256]",
-        "name" : "logits",
+        "shape" : "[1, 64, 16384]",
+        "name" : "logits_0",
+        "type" : "MultiArray"
+      },
+      {
+        "hasShapeFlexibility" : "0",
+        "isOptional" : "0",
+        "dataType" : "Float16",
+        "formattedType" : "MultiArray (Float16 1 × 64 × 16384)",
+        "shortDescription" : "",
+        "shape" : "[1, 64, 16384]",
+        "name" : "logits_1",
+        "type" : "MultiArray"
+      },
+      {
+        "hasShapeFlexibility" : "0",
+        "isOptional" : "0",
+        "dataType" : "Float16",
+        "formattedType" : "MultiArray (Float16 1 × 64 × 16384)",
+        "shortDescription" : "",
+        "shape" : "[1, 64, 16384]",
+        "name" : "logits_2",
+        "type" : "MultiArray"
+      },
+      {
+        "hasShapeFlexibility" : "0",
+        "isOptional" : "0",
+        "dataType" : "Float16",
+        "formattedType" : "MultiArray (Float16 1 × 64 × 16384)",
+        "shortDescription" : "",
+        "shape" : "[1, 64, 16384]",
+        "name" : "logits_3",
+        "type" : "MultiArray"
+      },
+      {
+        "hasShapeFlexibility" : "0",
+        "isOptional" : "0",
+        "dataType" : "Float16",
+        "formattedType" : "MultiArray (Float16 1 × 64 × 16384)",
+        "shortDescription" : "",
+        "shape" : "[1, 64, 16384]",
+        "name" : "logits_4",
+        "type" : "MultiArray"
+      },
+      {
+        "hasShapeFlexibility" : "0",
+        "isOptional" : "0",
+        "dataType" : "Float16",
+        "formattedType" : "MultiArray (Float16 1 × 64 × 16384)",
+        "shortDescription" : "",
+        "shape" : "[1, 64, 16384]",
+        "name" : "logits_5",
+        "type" : "MultiArray"
+      },
+      {
+        "hasShapeFlexibility" : "0",
+        "isOptional" : "0",
+        "dataType" : "Float16",
+        "formattedType" : "MultiArray (Float16 1 × 64 × 16384)",
+        "shortDescription" : "",
+        "shape" : "[1, 64, 16384]",
+        "name" : "logits_6",
+        "type" : "MultiArray"
+      },
+      {
+        "hasShapeFlexibility" : "0",
+        "isOptional" : "0",
+        "dataType" : "Float16",
+        "formattedType" : "MultiArray (Float16 1 × 64 × 13568)",
+        "shortDescription" : "",
+        "shape" : "[1, 64, 13568]",
+        "name" : "logits_7",
         "type" : "MultiArray"
       }
     ],
@@ -19,7 +89,7 @@
     ],
     "specificationVersion" : 7,
     "mlProgramOperationTypeHistogram" : {
-      "Concat" : 2,
+      "Concat" : 1,
       "Ios16.mul" : 2,
       "Squeeze" : 1,
       "Transpose" : 1,
@@ -58,7 +128,7 @@
         "type" : "MultiArray"
       }
     ],
-    "generatedClassName" : "Llama_3_2_1B_Instruct_2024_10_10_23_56_41_chunk6",
+    "generatedClassName" : "Llama_3_2_1B_Instruct_2024_10_13_15_34_32_chunk3",
     "method" : "predict"
   }
 ]

Llama-3.2-1B-Instruct_chunk6.mlmodelc/model.mil

@@ -27,51 +27,48 @@ program(1.0)
             tensor<fp16, [2048, 16384]> transpose_1_to_fp16 = const()[name = tensor<string, []>("transpose_1_to_fp16"), val = tensor<fp16, [2048, 16384]>(BLOBFILE(path = tensor<string, []>("@model_path/weights/weight.bin"), offset = tensor<uint64, []>(4416)))];
             tensor<fp16, [64, 16384]> matmul_0_cast_fp16 = matmul(transpose_x = matmul_0_transpose_x_0, transpose_y = matmul_0_transpose_y_0, x = reshape_0_cast_fp16, y = transpose_1_to_fp16)[name = tensor<string, []>("matmul_0_cast_fp16")];
             tensor<int32, [3]> concat_8 = const()[name = tensor<string, []>("concat_8"), val = tensor<int32, [3]>([1, 64, 16384])];
-            tensor<fp16, [1, 64, 16384]> reshape_2_cast_fp16 = reshape(shape = concat_8, x = matmul_0_cast_fp16)[name = tensor<string, []>("reshape_2_cast_fp16")];
+            tensor<fp16, [1, 64, 16384]> logits_0 = reshape(shape = concat_8, x = matmul_0_cast_fp16)[name = tensor<string, []>("reshape_2_cast_fp16")];
             tensor<bool, []> matmul_1_transpose_x_0 = const()[name = tensor<string, []>("matmul_1_transpose_x_0"), val = tensor<bool, []>(false)];
             tensor<bool, []> matmul_1_transpose_y_0 = const()[name = tensor<string, []>("matmul_1_transpose_y_0"), val = tensor<bool, []>(false)];
             tensor<fp16, [2048, 16384]> transpose_3_to_fp16 = const()[name = tensor<string, []>("transpose_3_to_fp16"), val = tensor<fp16, [2048, 16384]>(BLOBFILE(path = tensor<string, []>("@model_path/weights/weight.bin"), offset = tensor<uint64, []>(67113344)))];
             tensor<fp16, [64, 16384]> matmul_1_cast_fp16 = matmul(transpose_x = matmul_1_transpose_x_0, transpose_y = matmul_1_transpose_y_0, x = reshape_0_cast_fp16, y = transpose_3_to_fp16)[name = tensor<string, []>("matmul_1_cast_fp16")];
             tensor<int32, [3]> concat_16 = const()[name = tensor<string, []>("concat_16"), val = tensor<int32, [3]>([1, 64, 16384])];
-            tensor<fp16, [1, 64, 16384]> reshape_5_cast_fp16 = reshape(shape = concat_16, x = matmul_1_cast_fp16)[name = tensor<string, []>("reshape_5_cast_fp16")];
+            tensor<fp16, [1, 64, 16384]> logits_1 = reshape(shape = concat_16, x = matmul_1_cast_fp16)[name = tensor<string, []>("reshape_5_cast_fp16")];
             tensor<bool, []> matmul_2_transpose_x_0 = const()[name = tensor<string, []>("matmul_2_transpose_x_0"), val = tensor<bool, []>(false)];
             tensor<bool, []> matmul_2_transpose_y_0 = const()[name = tensor<string, []>("matmul_2_transpose_y_0"), val = tensor<bool, []>(false)];
             tensor<fp16, [2048, 16384]> transpose_5_to_fp16 = const()[name = tensor<string, []>("transpose_5_to_fp16"), val = tensor<fp16, [2048, 16384]>(BLOBFILE(path = tensor<string, []>("@model_path/weights/weight.bin"), offset = tensor<uint64, []>(134222272)))];
             tensor<fp16, [64, 16384]> matmul_2_cast_fp16 = matmul(transpose_x = matmul_2_transpose_x_0, transpose_y = matmul_2_transpose_y_0, x = reshape_0_cast_fp16, y = transpose_5_to_fp16)[name = tensor<string, []>("matmul_2_cast_fp16")];
             tensor<int32, [3]> concat_24 = const()[name = tensor<string, []>("concat_24"), val = tensor<int32, [3]>([1, 64, 16384])];
-            tensor<fp16, [1, 64, 16384]> reshape_8_cast_fp16 = reshape(shape = concat_24, x = matmul_2_cast_fp16)[name = tensor<string, []>("reshape_8_cast_fp16")];
+            tensor<fp16, [1, 64, 16384]> logits_2 = reshape(shape = concat_24, x = matmul_2_cast_fp16)[name = tensor<string, []>("reshape_8_cast_fp16")];
             tensor<bool, []> matmul_3_transpose_x_0 = const()[name = tensor<string, []>("matmul_3_transpose_x_0"), val = tensor<bool, []>(false)];
             tensor<bool, []> matmul_3_transpose_y_0 = const()[name = tensor<string, []>("matmul_3_transpose_y_0"), val = tensor<bool, []>(false)];
             tensor<fp16, [2048, 16384]> transpose_7_to_fp16 = const()[name = tensor<string, []>("transpose_7_to_fp16"), val = tensor<fp16, [2048, 16384]>(BLOBFILE(path = tensor<string, []>("@model_path/weights/weight.bin"), offset = tensor<uint64, []>(201331200)))];
             tensor<fp16, [64, 16384]> matmul_3_cast_fp16 = matmul(transpose_x = matmul_3_transpose_x_0, transpose_y = matmul_3_transpose_y_0, x = reshape_0_cast_fp16, y = transpose_7_to_fp16)[name = tensor<string, []>("matmul_3_cast_fp16")];
             tensor<int32, [3]> concat_32 = const()[name = tensor<string, []>("concat_32"), val = tensor<int32, [3]>([1, 64, 16384])];
-            tensor<fp16, [1, 64, 16384]> reshape_11_cast_fp16 = reshape(shape = concat_32, x = matmul_3_cast_fp16)[name = tensor<string, []>("reshape_11_cast_fp16")];
+            tensor<fp16, [1, 64, 16384]> logits_3 = reshape(shape = concat_32, x = matmul_3_cast_fp16)[name = tensor<string, []>("reshape_11_cast_fp16")];
             tensor<bool, []> matmul_4_transpose_x_0 = const()[name = tensor<string, []>("matmul_4_transpose_x_0"), val = tensor<bool, []>(false)];
             tensor<bool, []> matmul_4_transpose_y_0 = const()[name = tensor<string, []>("matmul_4_transpose_y_0"), val = tensor<bool, []>(false)];
             tensor<fp16, [2048, 16384]> transpose_9_to_fp16 = const()[name = tensor<string, []>("transpose_9_to_fp16"), val = tensor<fp16, [2048, 16384]>(BLOBFILE(path = tensor<string, []>("@model_path/weights/weight.bin"), offset = tensor<uint64, []>(268440128)))];
             tensor<fp16, [64, 16384]> matmul_4_cast_fp16 = matmul(transpose_x = matmul_4_transpose_x_0, transpose_y = matmul_4_transpose_y_0, x = reshape_0_cast_fp16, y = transpose_9_to_fp16)[name = tensor<string, []>("matmul_4_cast_fp16")];
             tensor<int32, [3]> concat_40 = const()[name = tensor<string, []>("concat_40"), val = tensor<int32, [3]>([1, 64, 16384])];
-            tensor<fp16, [1, 64, 16384]> reshape_14_cast_fp16 = reshape(shape = concat_40, x = matmul_4_cast_fp16)[name = tensor<string, []>("reshape_14_cast_fp16")];
+            tensor<fp16, [1, 64, 16384]> logits_4 = reshape(shape = concat_40, x = matmul_4_cast_fp16)[name = tensor<string, []>("reshape_14_cast_fp16")];
             tensor<bool, []> matmul_5_transpose_x_0 = const()[name = tensor<string, []>("matmul_5_transpose_x_0"), val = tensor<bool, []>(false)];
             tensor<bool, []> matmul_5_transpose_y_0 = const()[name = tensor<string, []>("matmul_5_transpose_y_0"), val = tensor<bool, []>(false)];
             tensor<fp16, [2048, 16384]> transpose_11_to_fp16 = const()[name = tensor<string, []>("transpose_11_to_fp16"), val = tensor<fp16, [2048, 16384]>(BLOBFILE(path = tensor<string, []>("@model_path/weights/weight.bin"), offset = tensor<uint64, []>(335549056)))];
             tensor<fp16, [64, 16384]> matmul_5_cast_fp16 = matmul(transpose_x = matmul_5_transpose_x_0, transpose_y = matmul_5_transpose_y_0, x = reshape_0_cast_fp16, y = transpose_11_to_fp16)[name = tensor<string, []>("matmul_5_cast_fp16")];
             tensor<int32, [3]> concat_48 = const()[name = tensor<string, []>("concat_48"), val = tensor<int32, [3]>([1, 64, 16384])];
-            tensor<fp16, [1, 64, 16384]> reshape_17_cast_fp16 = reshape(shape = concat_48, x = matmul_5_cast_fp16)[name = tensor<string, []>("reshape_17_cast_fp16")];
+            tensor<fp16, [1, 64, 16384]> logits_5 = reshape(shape = concat_48, x = matmul_5_cast_fp16)[name = tensor<string, []>("reshape_17_cast_fp16")];
             tensor<bool, []> matmul_6_transpose_x_0 = const()[name = tensor<string, []>("matmul_6_transpose_x_0"), val = tensor<bool, []>(false)];
             tensor<bool, []> matmul_6_transpose_y_0 = const()[name = tensor<string, []>("matmul_6_transpose_y_0"), val = tensor<bool, []>(false)];
             tensor<fp16, [2048, 16384]> transpose_13_to_fp16 = const()[name = tensor<string, []>("transpose_13_to_fp16"), val = tensor<fp16, [2048, 16384]>(BLOBFILE(path = tensor<string, []>("@model_path/weights/weight.bin"), offset = tensor<uint64, []>(402657984)))];
             tensor<fp16, [64, 16384]> matmul_6_cast_fp16 = matmul(transpose_x = matmul_6_transpose_x_0, transpose_y = matmul_6_transpose_y_0, x = reshape_0_cast_fp16, y = transpose_13_to_fp16)[name = tensor<string, []>("matmul_6_cast_fp16")];
             tensor<int32, [3]> concat_56 = const()[name = tensor<string, []>("concat_56"), val = tensor<int32, [3]>([1, 64, 16384])];
-            tensor<fp16, [1, 64, 16384]> reshape_20_cast_fp16 = reshape(shape = concat_56, x = matmul_6_cast_fp16)[name = tensor<string, []>("reshape_20_cast_fp16")];
+            tensor<fp16, [1, 64, 16384]> logits_6 = reshape(shape = concat_56, x = matmul_6_cast_fp16)[name = tensor<string, []>("reshape_20_cast_fp16")];
             tensor<bool, []> matmul_7_transpose_x_0 = const()[name = tensor<string, []>("matmul_7_transpose_x_0"), val = tensor<bool, []>(false)];
             tensor<bool, []> matmul_7_transpose_y_0 = const()[name = tensor<string, []>("matmul_7_transpose_y_0"), val = tensor<bool, []>(false)];
             tensor<fp16, [2048, 13568]> transpose_15_to_fp16 = const()[name = tensor<string, []>("transpose_15_to_fp16"), val = tensor<fp16, [2048, 13568]>(BLOBFILE(path = tensor<string, []>("@model_path/weights/weight.bin"), offset = tensor<uint64, []>(469766912)))];
             tensor<fp16, [64, 13568]> matmul_7_cast_fp16 = matmul(transpose_x = matmul_7_transpose_x_0, transpose_y = matmul_7_transpose_y_0, x = reshape_0_cast_fp16, y = transpose_15_to_fp16)[name = tensor<string, []>("matmul_7_cast_fp16")];
             tensor<int32, [3]> concat_64 = const()[name = tensor<string, []>("concat_64"), val = tensor<int32, [3]>([1, 64, 13568])];
-            tensor<fp16, [1, 64, 13568]> reshape_23_cast_fp16 = reshape(shape = concat_64, x = matmul_7_cast_fp16)[name = tensor<string, []>("reshape_23_cast_fp16")];
-            tensor<int32, []> var_99 = const()[name = tensor<string, []>("op_99"), val = tensor<int32, []>(-1)];
-            tensor<bool, []> var_100_interleave_0 = const()[name = tensor<string, []>("op_100_interleave_0"), val = tensor<bool, []>(false)];
-            tensor<fp16, [1, 64, 128256]> logits = concat(axis = var_99, interleave = var_100_interleave_0, values = (reshape_2_cast_fp16, reshape_5_cast_fp16, reshape_8_cast_fp16, reshape_11_cast_fp16, reshape_14_cast_fp16, reshape_17_cast_fp16, reshape_20_cast_fp16, reshape_23_cast_fp16))[name = tensor<string, []>("op_100_cast_fp16")];
-        } -> (logits);
+            tensor<fp16, [1, 64, 13568]> logits_7 = reshape(shape = concat_64, x = matmul_7_cast_fp16)[name = tensor<string, []>("reshape_23_cast_fp16")];
+        } -> (logits_0, logits_1, logits_2, logits_3, logits_4, logits_5, logits_6, logits_7);
 }

logit-processor.mlmodelc/analytics/coremldata.bin

@@ -1,3 +1,3 @@
 version https://git-lfs.github.com/spec/v1
-oid sha256:0ad03dc247f59282bf008d857db8620b0ad600eb939bfa2a4e8a78438e1c2573
+oid sha256:0cea8f79e82c95d93f772797047802fb88c7fc82dfcef790e69a2f274a104623
 size 243

logit-processor.mlmodelc/coremldata.bin

@@ -1,3 +1,3 @@
 version https://git-lfs.github.com/spec/v1
-oid sha256:ccca55190c5da56bfc175471f3239eeeb7bffece8d38d565de9443edef9c9148
-size 378
+oid sha256:9d70f289c1552a24ba6ca721405ea0653ac7125a91297ea3d32b30363c2afd3c
+size 503

logit-processor.mlmodelc/metadata.json

@@ -6,9 +6,9 @@
         "hasShapeFlexibility" : "0",
         "isOptional" : "0",
         "dataType" : "Int32",
-        "formattedType" : "MultiArray (Int32)",
+        "formattedType" : "MultiArray (Int32 1 × 64)",
         "shortDescription" : "",
-        "shape" : "[]",
+        "shape" : "[1, 64]",
         "name" : "argmax",
         "type" : "MultiArray"
       }
@@ -18,7 +18,11 @@
     ],
     "specificationVersion" : 7,
     "mlProgramOperationTypeHistogram" : {
-      "Ios16.reduceArgmax" : 1
+      "Ios16.add" : 7,
+      "Ios16.topk" : 9,
+      "Ios16.gatherAlongAxis" : 1,
+      "Concat" : 2,
+      "Squeeze" : 1
     },
     "computePrecision" : "Mixed (Float16, Int32)",
     "isUpdatable" : "0",
@@ -40,19 +44,87 @@
     },
     "inputSchema" : [
       {
+        "hasShapeFlexibility" : "0",
+        "isOptional" : "0",
+        "dataType" : "Float16",
+        "formattedType" : "MultiArray (Float16 1 × 64 × 16384)",
+        "shortDescription" : "",
+        "shape" : "[1, 64, 16384]",
+        "name" : "logits_0",
+        "type" : "MultiArray"
+      },
+      {
+        "hasShapeFlexibility" : "0",
+        "isOptional" : "0",
+        "dataType" : "Float16",
+        "formattedType" : "MultiArray (Float16 1 × 64 × 16384)",
+        "shortDescription" : "",
+        "shape" : "[1, 64, 16384]",
+        "name" : "logits_1",
+        "type" : "MultiArray"
+      },
+      {
+        "hasShapeFlexibility" : "0",
+        "isOptional" : "0",
+        "dataType" : "Float16",
+        "formattedType" : "MultiArray (Float16 1 × 64 × 16384)",
+        "shortDescription" : "",
+        "shape" : "[1, 64, 16384]",
+        "name" : "logits_2",
+        "type" : "MultiArray"
+      },
+      {
+        "hasShapeFlexibility" : "0",
+        "isOptional" : "0",
+        "dataType" : "Float16",
+        "formattedType" : "MultiArray (Float16 1 × 64 × 16384)",
+        "shortDescription" : "",
+        "shape" : "[1, 64, 16384]",
+        "name" : "logits_3",
+        "type" : "MultiArray"
+      },
+      {
+        "hasShapeFlexibility" : "0",
+        "isOptional" : "0",
+        "dataType" : "Float16",
+        "formattedType" : "MultiArray (Float16 1 × 64 × 16384)",
+        "shortDescription" : "",
+        "shape" : "[1, 64, 16384]",
+        "name" : "logits_4",
+        "type" : "MultiArray"
+      },
+      {
+        "hasShapeFlexibility" : "0",
+        "isOptional" : "0",
+        "dataType" : "Float16",
+        "formattedType" : "MultiArray (Float16 1 × 64 × 16384)",
         "shortDescription" : "",
+        "shape" : "[1, 64, 16384]",
+        "name" : "logits_5",
+        "type" : "MultiArray"
+      },
+      {
+        "hasShapeFlexibility" : "0",
+        "isOptional" : "0",
         "dataType" : "Float16",
-        "hasShapeFlexibility" : "1",
+        "formattedType" : "MultiArray (Float16 1 × 64 × 16384)",
+        "shortDescription" : "",
+        "shape" : "[1, 64, 16384]",
+        "name" : "logits_6",
+        "type" : "MultiArray"
+      },
+      {
+        "hasShapeFlexibility" : "0",
         "isOptional" : "0",
-        "shapeFlexibility" : "1 × 511 × 32000 | 1 × 1 × 32000 | 1 × 2 × 32000 | 1 × 4 × 32000 | 1 × 64 × 32000 | 1 × 64 × 128256 | 1 × 512 × 32000",
-        "formattedType" : "MultiArray (Float16 1 × 511 × 32000)",
-        "type" : "MultiArray",
-        "shape" : "[1, 511, 32000]",
-        "name" : "logits",
-        "enumeratedShapes" : "[[1, 511, 32000], [1, 1, 32000], [1, 2, 32000], [1, 4, 32000], [1, 64, 32000], [1, 64, 128256], [1, 512, 32000]]"
+        "dataType" : "Float16",
+        "formattedType" : "MultiArray (Float16 1 × 64 × 13568)",
+        "shortDescription" : "",
+        "shape" : "[1, 64, 13568]",
+        "name" : "logits_7",
+        "type" : "MultiArray"
       }
     ],
-    "generatedClassName" : "logit_processor",
+    "generatedClassName" : "split_logit_processor",
     "method" : "predict"
   }
 ]

logit-processor.mlmodelc/model.mil

@@ -1,9 +1,84 @@
 program(1.0)
 [buildInfo = dict<tensor<string, []>, tensor<string, []>>({{"coremlc-component-MIL", "3304.5.2"}, {"coremlc-version", "3304.6.2"}, {"coremltools-component-torch", "2.1.0"}, {"coremltools-source-dialect", "TorchScript"}, {"coremltools-version", "8.0b1"}})]
 {
-    func main<ios16>(tensor<fp16, [1, ?, ?]> logits) [FlexibleShapeInformation = tuple<tuple<tensor<string, []>, dict<tensor<string, []>, tensor<int32, [?]>>>, tuple<tensor<string, []>, dict<tensor<string, []>, dict<tensor<string, []>, tensor<int32, [?]>>>>>((("DefaultShapes", {{"logits", [1, 511, 32000]}}), ("EnumeratedShapes", {{"logits_1_1_1_1_32000_", {{"logits", [1, 1, 32000]}}}, {"logits_1_1_1_2_32000_", {{"logits", [1, 2, 32000]}}}, {"logits_1_1_1_4_32000_", {{"logits", [1, 4, 32000]}}}, {"logits_1_1_1_511_32000_", {{"logits", [1, 511, 32000]}}}, {"logits_1_1_1_512_32000_", {{"logits", [1, 512, 32000]}}}, {"logits_1_1_1_64_128256_", {{"logits", [1, 64, 128256]}}}, {"logits_1_1_1_64_32000_", {{"logits", [1, 64, 32000]}}}})))] {
-            tensor<int32, []> var_2 = const()[name = tensor<string, []>("op_2"), val = tensor<int32, []>(-1)];
-            tensor<bool, []> var_3 = const()[name = tensor<string, []>("op_3"), val = tensor<bool, []>(false)];
-            tensor<int32, [1, ?]> argmax = reduce_argmax(axis = var_2, keep_dims = var_3, x = logits)[name = tensor<string, []>("op_4_cast_fp16")];
+    func main<ios16>(tensor<fp16, [1, 64, 16384]> logits_0, tensor<fp16, [1, 64, 16384]> logits_1, tensor<fp16, [1, 64, 16384]> logits_2, tensor<fp16, [1, 64, 16384]> logits_3, tensor<fp16, [1, 64, 16384]> logits_4, tensor<fp16, [1, 64, 16384]> logits_5, tensor<fp16, [1, 64, 16384]> logits_6, tensor<fp16, [1, 64, 13568]> logits_7) {
+            tensor<int32, [1]> chunk_size = const()[name = tensor<string, []>("chunk_size"), val = tensor<int32, [1]>([16384])];
+            tensor<int32, []> var_12 = const()[name = tensor<string, []>("op_12"), val = tensor<int32, []>(1)];
+            tensor<int32, []> var_16_axis_0 = const()[name = tensor<string, []>("op_16_axis_0"), val = tensor<int32, []>(-1)];
+            tensor<bool, []> var_16_ascending_0 = const()[name = tensor<string, []>("op_16_ascending_0"), val = tensor<bool, []>(false)];
+            tensor<bool, []> var_16_sort_0 = const()[name = tensor<string, []>("op_16_sort_0"), val = tensor<bool, []>(false)];
+            tensor<bool, []> var_16_return_indices_0 = const()[name = tensor<string, []>("op_16_return_indices_0"), val = tensor<bool, []>(true)];
+            tensor<fp16, [1, 64, 1]> var_16_cast_fp16_0, tensor<int32, [1, 64, 1]> var_16_cast_fp16_1 = topk(ascending = var_16_ascending_0, axis = var_16_axis_0, k = var_12, return_indices = var_16_return_indices_0, sort = var_16_sort_0, x = logits_0)[name = tensor<string, []>("op_16_cast_fp16")];
+            tensor<int32, []> var_22 = const()[name = tensor<string, []>("op_22"), val = tensor<int32, []>(1)];
+            tensor<int32, []> var_26_axis_0 = const()[name = tensor<string, []>("op_26_axis_0"), val = tensor<int32, []>(-1)];
+            tensor<bool, []> var_26_ascending_0 = const()[name = tensor<string, []>("op_26_ascending_0"), val = tensor<bool, []>(false)];
+            tensor<bool, []> var_26_sort_0 = const()[name = tensor<string, []>("op_26_sort_0"), val = tensor<bool, []>(false)];
+            tensor<bool, []> var_26_return_indices_0 = const()[name = tensor<string, []>("op_26_return_indices_0"), val = tensor<bool, []>(true)];
+            tensor<fp16, [1, 64, 1]> var_26_cast_fp16_0, tensor<int32, [1, 64, 1]> var_26_cast_fp16_1 = topk(ascending = var_26_ascending_0, axis = var_26_axis_0, k = var_22, return_indices = var_26_return_indices_0, sort = var_26_sort_0, x = logits_1)[name = tensor<string, []>("op_26_cast_fp16")];
+            tensor<int32, [1, 64, 1]> var_31 = add(x = var_26_cast_fp16_1, y = chunk_size)[name = tensor<string, []>("op_31")];
+            tensor<int32, []> var_32 = const()[name = tensor<string, []>("op_32"), val = tensor<int32, []>(1)];
+            tensor<int32, []> var_36_axis_0 = const()[name = tensor<string, []>("op_36_axis_0"), val = tensor<int32, []>(-1)];
+            tensor<bool, []> var_36_ascending_0 = const()[name = tensor<string, []>("op_36_ascending_0"), val = tensor<bool, []>(false)];
+            tensor<bool, []> var_36_sort_0 = const()[name = tensor<string, []>("op_36_sort_0"), val = tensor<bool, []>(false)];
+            tensor<bool, []> var_36_return_indices_0 = const()[name = tensor<string, []>("op_36_return_indices_0"), val = tensor<bool, []>(true)];
+            tensor<fp16, [1, 64, 1]> var_36_cast_fp16_0, tensor<int32, [1, 64, 1]> var_36_cast_fp16_1 = topk(ascending = var_36_ascending_0, axis = var_36_axis_0, k = var_32, return_indices = var_36_return_indices_0, sort = var_36_sort_0, x = logits_2)[name = tensor<string, []>("op_36_cast_fp16")];
+            tensor<int32, [1]> var_39 = const()[name = tensor<string, []>("op_39"), val = tensor<int32, [1]>([32768])];
+            tensor<int32, [1, 64, 1]> var_41 = add(x = var_36_cast_fp16_1, y = var_39)[name = tensor<string, []>("op_41")];
+            tensor<int32, []> var_42 = const()[name = tensor<string, []>("op_42"), val = tensor<int32, []>(1)];
+            tensor<int32, []> var_46_axis_0 = const()[name = tensor<string, []>("op_46_axis_0"), val = tensor<int32, []>(-1)];
+            tensor<bool, []> var_46_ascending_0 = const()[name = tensor<string, []>("op_46_ascending_0"), val = tensor<bool, []>(false)];
+            tensor<bool, []> var_46_sort_0 = const()[name = tensor<string, []>("op_46_sort_0"), val = tensor<bool, []>(false)];
+            tensor<bool, []> var_46_return_indices_0 = const()[name = tensor<string, []>("op_46_return_indices_0"), val = tensor<bool, []>(true)];
+            tensor<fp16, [1, 64, 1]> var_46_cast_fp16_0, tensor<int32, [1, 64, 1]> var_46_cast_fp16_1 = topk(ascending = var_46_ascending_0, axis = var_46_axis_0, k = var_42, return_indices = var_46_return_indices_0, sort = var_46_sort_0, x = logits_3)[name = tensor<string, []>("op_46_cast_fp16")];
+            tensor<int32, [1]> var_49 = const()[name = tensor<string, []>("op_49"), val = tensor<int32, [1]>([49152])];
+            tensor<int32, [1, 64, 1]> var_51 = add(x = var_46_cast_fp16_1, y = var_49)[name = tensor<string, []>("op_51")];
+            tensor<int32, []> var_52 = const()[name = tensor<string, []>("op_52"), val = tensor<int32, []>(1)];
+            tensor<int32, []> var_56_axis_0 = const()[name = tensor<string, []>("op_56_axis_0"), val = tensor<int32, []>(-1)];
+            tensor<bool, []> var_56_ascending_0 = const()[name = tensor<string, []>("op_56_ascending_0"), val = tensor<bool, []>(false)];
+            tensor<bool, []> var_56_sort_0 = const()[name = tensor<string, []>("op_56_sort_0"), val = tensor<bool, []>(false)];
+            tensor<bool, []> var_56_return_indices_0 = const()[name = tensor<string, []>("op_56_return_indices_0"), val = tensor<bool, []>(true)];
+            tensor<fp16, [1, 64, 1]> var_56_cast_fp16_0, tensor<int32, [1, 64, 1]> var_56_cast_fp16_1 = topk(ascending = var_56_ascending_0, axis = var_56_axis_0, k = var_52, return_indices = var_56_return_indices_0, sort = var_56_sort_0, x = logits_4)[name = tensor<string, []>("op_56_cast_fp16")];
+            tensor<int32, [1]> var_59 = const()[name = tensor<string, []>("op_59"), val = tensor<int32, [1]>([65536])];
+            tensor<int32, [1, 64, 1]> var_61 = add(x = var_56_cast_fp16_1, y = var_59)[name = tensor<string, []>("op_61")];
+            tensor<int32, []> var_62 = const()[name = tensor<string, []>("op_62"), val = tensor<int32, []>(1)];
+            tensor<int32, []> var_66_axis_0 = const()[name = tensor<string, []>("op_66_axis_0"), val = tensor<int32, []>(-1)];
+            tensor<bool, []> var_66_ascending_0 = const()[name = tensor<string, []>("op_66_ascending_0"), val = tensor<bool, []>(false)];
+            tensor<bool, []> var_66_sort_0 = const()[name = tensor<string, []>("op_66_sort_0"), val = tensor<bool, []>(false)];
+            tensor<bool, []> var_66_return_indices_0 = const()[name = tensor<string, []>("op_66_return_indices_0"), val = tensor<bool, []>(true)];
+            tensor<fp16, [1, 64, 1]> var_66_cast_fp16_0, tensor<int32, [1, 64, 1]> var_66_cast_fp16_1 = topk(ascending = var_66_ascending_0, axis = var_66_axis_0, k = var_62, return_indices = var_66_return_indices_0, sort = var_66_sort_0, x = logits_5)[name = tensor<string, []>("op_66_cast_fp16")];
+            tensor<int32, [1]> var_69 = const()[name = tensor<string, []>("op_69"), val = tensor<int32, [1]>([81920])];
+            tensor<int32, [1, 64, 1]> var_71 = add(x = var_66_cast_fp16_1, y = var_69)[name = tensor<string, []>("op_71")];
+            tensor<int32, []> var_72 = const()[name = tensor<string, []>("op_72"), val = tensor<int32, []>(1)];
+            tensor<int32, []> var_76_axis_0 = const()[name = tensor<string, []>("op_76_axis_0"), val = tensor<int32, []>(-1)];
+            tensor<bool, []> var_76_ascending_0 = const()[name = tensor<string, []>("op_76_ascending_0"), val = tensor<bool, []>(false)];
+            tensor<bool, []> var_76_sort_0 = const()[name = tensor<string, []>("op_76_sort_0"), val = tensor<bool, []>(false)];
+            tensor<bool, []> var_76_return_indices_0 = const()[name = tensor<string, []>("op_76_return_indices_0"), val = tensor<bool, []>(true)];
+            tensor<fp16, [1, 64, 1]> var_76_cast_fp16_0, tensor<int32, [1, 64, 1]> var_76_cast_fp16_1 = topk(ascending = var_76_ascending_0, axis = var_76_axis_0, k = var_72, return_indices = var_76_return_indices_0, sort = var_76_sort_0, x = logits_6)[name = tensor<string, []>("op_76_cast_fp16")];
+            tensor<int32, [1]> var_79 = const()[name = tensor<string, []>("op_79"), val = tensor<int32, [1]>([98304])];
+            tensor<int32, [1, 64, 1]> var_81 = add(x = var_76_cast_fp16_1, y = var_79)[name = tensor<string, []>("op_81")];
+            tensor<int32, []> var_82 = const()[name = tensor<string, []>("op_82"), val = tensor<int32, []>(1)];
+            tensor<int32, []> cv_axis_0 = const()[name = tensor<string, []>("cv_axis_0"), val = tensor<int32, []>(-1)];
+            tensor<bool, []> cv_ascending_0 = const()[name = tensor<string, []>("cv_ascending_0"), val = tensor<bool, []>(false)];
+            tensor<bool, []> cv_sort_0 = const()[name = tensor<string, []>("cv_sort_0"), val = tensor<bool, []>(false)];
+            tensor<bool, []> cv_return_indices_0 = const()[name = tensor<string, []>("cv_return_indices_0"), val = tensor<bool, []>(true)];
+            tensor<fp16, [1, 64, 1]> cv_cast_fp16_0, tensor<int32, [1, 64, 1]> cv_cast_fp16_1 = topk(ascending = cv_ascending_0, axis = cv_axis_0, k = var_82, return_indices = cv_return_indices_0, sort = cv_sort_0, x = logits_7)[name = tensor<string, []>("cv_cast_fp16")];
+            tensor<int32, [1]> var_89 = const()[name = tensor<string, []>("op_89"), val = tensor<int32, [1]>([114688])];
+            tensor<int32, [1, 64, 1]> var_91 = add(x = cv_cast_fp16_1, y = var_89)[name = tensor<string, []>("op_91")];
+            tensor<int32, []> var_93 = const()[name = tensor<string, []>("op_93"), val = tensor<int32, []>(-1)];
+            tensor<bool, []> values_interleave_0 = const()[name = tensor<string, []>("values_interleave_0"), val = tensor<bool, []>(false)];
+            tensor<fp16, [1, 64, 8]> values_cast_fp16 = concat(axis = var_93, interleave = values_interleave_0, values = (var_16_cast_fp16_0, var_26_cast_fp16_0, var_36_cast_fp16_0, var_46_cast_fp16_0, var_56_cast_fp16_0, var_66_cast_fp16_0, var_76_cast_fp16_0, cv_cast_fp16_0))[name = tensor<string, []>("values_cast_fp16")];
+            tensor<int32, []> var_96 = const()[name = tensor<string, []>("op_96"), val = tensor<int32, []>(-1)];
+            tensor<bool, []> indices_interleave_0 = const()[name = tensor<string, []>("indices_interleave_0"), val = tensor<bool, []>(false)];
+            tensor<int32, [1, 64, 8]> indices = concat(axis = var_96, interleave = indices_interleave_0, values = (var_16_cast_fp16_1, var_31, var_41, var_51, var_61, var_71, var_81, var_91))[name = tensor<string, []>("indices")];
+            tensor<int32, []> var_98 = const()[name = tensor<string, []>("op_98"), val = tensor<int32, []>(1)];
+            tensor<int32, []> var_102_axis_0 = const()[name = tensor<string, []>("op_102_axis_0"), val = tensor<int32, []>(-1)];
+            tensor<bool, []> var_102_ascending_0 = const()[name = tensor<string, []>("op_102_ascending_0"), val = tensor<bool, []>(false)];
+            tensor<bool, []> var_102_sort_0 = const()[name = tensor<string, []>("op_102_sort_0"), val = tensor<bool, []>(true)];
+            tensor<bool, []> var_102_return_indices_0 = const()[name = tensor<string, []>("op_102_return_indices_0"), val = tensor<bool, []>(true)];
+            tensor<fp16, [1, 64, 1]> var_102_cast_fp16_0, tensor<int32, [1, 64, 1]> var_102_cast_fp16_1 = topk(ascending = var_102_ascending_0, axis = var_102_axis_0, k = var_98, return_indices = var_102_return_indices_0, sort = var_102_sort_0, x = values_cast_fp16)[name = tensor<string, []>("op_102_cast_fp16")];
+            tensor<int32, []> var_104 = const()[name = tensor<string, []>("op_104"), val = tensor<int32, []>(-1)];
+            tensor<int32, [1, 64, 1]> var_106 = gather_along_axis(axis = var_104, indices = var_102_cast_fp16_1, x = indices)[name = tensor<string, []>("op_106")];
+            tensor<int32, [1]> var_108_axes_0 = const()[name = tensor<string, []>("op_108_axes_0"), val = tensor<int32, [1]>([-1])];
+            tensor<int32, [1, 64]> argmax = squeeze(axes = var_108_axes_0, x = var_106)[name = tensor<string, []>("op_108")];
         } -> (argmax);
 }