From d80341831aec02598bd15904cb35f24ba1d4edde Mon Sep 17 00:00:00 2001 From: Yuanqiang Liu Date: Sat, 27 Jul 2024 18:53:08 +0800 Subject: [PATCH] update --- .../mlir_tests/cpu_ops/custom_call_byteir_arg_max.mlir | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/tests/numerical_test/mlir_tests/cpu_ops/custom_call_byteir_arg_max.mlir b/tests/numerical_test/mlir_tests/cpu_ops/custom_call_byteir_arg_max.mlir index a9fdd7c45..beb6b07a7 100644 --- a/tests/numerical_test/mlir_tests/cpu_ops/custom_call_byteir_arg_max.mlir +++ b/tests/numerical_test/mlir_tests/cpu_ops/custom_call_byteir_arg_max.mlir @@ -1,4 +1,4 @@ -func.func @byteir.arg_max$return_2(%arg0: tensor<3x4xf32>) -> (tensor<3xf32>, tensor<3xi64>) { - %0:2 = stablehlo.custom_call @byteir.arg_max(%arg0) {byteir_attrs = {axis = 1 : i64, keep_dims = false, select_last_index = false}} : (tensor<3x4xf32>) -> (tensor<3xf32>, tensor<3xi64>) +func.func @byteir.arg_max$return_2(%arg0: tensor<3x128xf32>) -> (tensor<3xf32>, tensor<3xi64>) { + %0:2 = stablehlo.custom_call @byteir.arg_max(%arg0) {byteir_attrs = {axis = 1 : i64, keep_dims = false, select_last_index = false}} : (tensor<3x128xf32>) -> (tensor<3xf32>, tensor<3xi64>) return %0#0, %0#1 : tensor<3xf32>, tensor<3xi64> }