From f0b484c0bc59f0f39a93e111cee428fb6f596ccf Mon Sep 17 00:00:00 2001 From: Doru Bercea Date: Tue, 14 Jan 2020 10:37:05 -0500 Subject: [PATCH] Add test for transpose with permutation. --- test/mlir/onnx/onnx_shape_inference.mlir | 12 +++++++++++- 1 file changed, 11 insertions(+), 1 deletion(-) diff --git a/test/mlir/onnx/onnx_shape_inference.mlir b/test/mlir/onnx/onnx_shape_inference.mlir index 4cb9bec..aaa08a7 100644 --- a/test/mlir/onnx/onnx_shape_inference.mlir +++ b/test/mlir/onnx/onnx_shape_inference.mlir @@ -9,4 +9,14 @@ func @test_default_transpose(%arg0 : tensor<5x5x1x32xf32>) -> tensor<*xf32> { // CHECK-LABEL: test_default_transpose // CHECK: [[RES:%.+]] = "onnx.Transpose"(%arg0) : (tensor<5x5x1x32xf32>) -> tensor<32x1x5x5xf32> -// CHECK: return [[RES]] : tensor<32x1x5x5xf32> \ No newline at end of file +// CHECK: return [[RES]] : tensor<32x1x5x5xf32> + +/// Test shape inference for transposition when perm attribute is specified. +func @test_transpose(%arg0 : tensor<5x5x1x32xf32>) -> tensor<*xf32> { + %0 = "onnx.Transpose"(%arg0) {perm = [2, 0, 3, 1]} : (tensor<5x5x1x32xf32>) -> tensor<*xf32> + "std.return"(%0) : (tensor<*xf32>) -> () +} + +// CHECK-LABEL: test_transpose +// CHECK: [[RES_ATTR:%.+]] = "onnx.Transpose"(%arg0) {perm = [2, 0, 3, 1]} : (tensor<5x5x1x32xf32>) -> tensor<1x5x32x5xf32> +// CHECK: return [[RES_ATTR]] : tensor<1x5x32x5xf32> \ No newline at end of file