12 lines · plain
1// RUN: mlir-opt %s | mlir-opt | FileCheck %s2// RUN: mlir-opt %s --mlir-print-op-generic | mlir-opt | FileCheck %s3 4 5// -----6// Uses argmax as canonical example to validate constrained TOSA tensor shapes.7// CHECK-LABEL: argmax8func.func @test_argmax(%arg0: tensor<?xf32>) -> tensor<i32> {9 %0 = "tosa.argmax"(%arg0) {axis = 0 : i32} : (tensor<?xf32>) -> tensor<i32>10 return %0 : tensor<i32>11}12