xref: /llvm-project/mlir/test/Dialect/Tosa/constrained_shapes.mlir (revision 414709ee5c74ae35ac5eb204007fef35cc1c8079) 
1// RUN: mlir-opt %s | mlir-opt | FileCheck %s
2// RUN: mlir-opt %s --mlir-print-op-generic | mlir-opt | FileCheck %s
3
4
5// -----
6// Uses argmax as canonical example to validate constrained TOSA tensor shapes.
7// CHECK-LABEL: argmax
8func.func @test_argmax(%arg0: tensor<?xf32>) -> tensor<?xi32> {
9  %0 = "tosa.argmax"(%arg0) {axis = 0 : i32} : (tensor<?xf32>) -> tensor<?xi32>
10  return %0 : tensor<?xi32>
11}
12