// Returns a ShapedType for a permutation and the shape of input after
// applying the permutation to the given shape through a transpose.
class GetTransposedType<string perm> : NativeCodeCall<
  "GetTransposedType($0, " # perm # ")">;

// Function to map final permutation to initial permutation
// initial -> permutation1 -> permutation2 -> final

  return shaped_type.clone(new_shape);
}

// Returns a ShapedType for a permutation and the shape of input after
// applying the permutation to the given shape through a transpose.
inline ShapedType GetTransposedType(Value input,
                                    llvm::ArrayRef<int64_t> permutation_array) {
  auto input_type = input.getType().cast<ShapedType>();
  if (permutation_array.size() != input_type.getRank()) {

              /*stride_h=*/ExtractI32At<2>:$strides,
              /*stride_w=*/ExtractI32At<3>:$strides,
              /*fused_activation_function=*/TFL_AF_None,
              (returnType (GetTransposedType<"{0, 2, 3, 1}"> $src_op))),
              (Arith_ConstantOp
                  ConstantAttr<I32VectorElementsAttr<4>, "{0, 3, 1, 2}">))>;

def LegalizeArgMax : Pat<(TF_ArgMaxOp $input, $dim),