// RUN: tac-opt-all-backends -tfl-device-transform-gpu %s -split-input-file -verify-diagnostics | FileCheck %s

func.func @pack(%arg0: tensor<1xf32>, %arg1: tensor<1xf32>) -> tensor<2x1xf32> {
  %0 = "tfl.pack"(%arg0, %arg1) {axis = 0 : i32, values_count = 2 : i32} : (tensor<1xf32>, tensor<1xf32>) -> tensor<2x1xf32>
  func.return %0 : tensor<2x1xf32>
}

// CHECK:   func @pack(%[[VAL_0:.*]]: tensor<1xf32>, %[[VAL_1:.*]]: tensor<1xf32>) -> tensor<2x1xf32> {

def LegalizeNeg : Pat<(TF_NegOp $arg), (TFL_NegOp $arg)>;
def LegalizeOneHot : Pat<
  (TF_OneHotOp $indices, $depth, $on_value, $off_value, $axis),
  (TFL_OneHotOp $indices, $depth, $on_value, $off_value,
  (convertIntAttrTo32Bit $axis))>;
def LegalizePow : Pat<(TF_PowOp $x, $y), (TFL_PowOp $x, $y)>;
def LegalizeRange : Pat<(TF_RangeOp $start, $limit, $delta),

    int64_t axis = axis_int.getSExtValue();
    if (axis < 0) {
      axis += input_type.getRank();
    }
    if (axis < 0 || (input_type.hasRank() && axis >= input_type.getRank())) {
      return op.emitOpError(
          llvm::formatv("perm[{0}] must be in [-rank, rank)", index));
    }
    if (std::count(axes.begin(), axes.end(), axis) > 0) {
      return op.emitOpError(

      int64_t axis = axis_int.getSExtValue();
      if (axis < 0) {
        axis += input_type.getRank();
      }
      if (axis < 0 || (input_type.hasRank() && axis >= input_type.getRank())) {
        return transpose_op.emitOpError("perm must be in [-rank, rank)");
      }
      if (std::count(axes.begin(), axes.end(), axis) > 0) {

  // CHECK-SAME: ([[INPUT:%.*]]: tensor<1x128x128xf32>, [[FILTER:%.*]]: tensor<5x5x1x1xf32>, [[BIAS:%.*]]: tensor<128xf32>)
  // CHECK-NEXT: [[AXIS:%.*]] = "tf.Const"() <{value = dense<3> : tensor<i32>}> : () -> tensor<i32>
  // CHECK-NEXT: [[EXPAND:%.*]] = "tf.ExpandDims"([[INPUT]], [[AXIS]]) : (tensor<1x128x128xf32>, tensor<i32>) -> tensor<1x128x128x1xf32>

    %6 = tfl.add %arg1, %cst_5 {fused_activation_function = "NONE"} : tensor<i32>
    %7 = "tfl.gather"(%0, %arg2) {axis = 0 : i32} : (tensor<4x4x3xf32>, tensor<i32>) -> tensor<4x3xf32>
    %8 = "tfl.concatenation"(%7, %arg5) {axis = 1 : i32, fused_activation_function = "NONE"} : (tensor<4x3xf32>, tensor<4x2xf32>) -> tensor<4x5xf32>

                                << -rank << ", " << rank << ")";
      if (axis < 0) axis += rank;
    }

    if (batch_dims >= 0 && axis >= 0 && axis < batch_dims) {
      return op.emitOpError() << "requires axis (" << axis
                              << ") to be greater than or equal to batch_dims ("
                              << batch_dims << ")";
    }
  }
  return success();
}

// CHECK-NEXT:      %17 = "tfl.gather"(%14, %16) <{axis = 0 : i32, batch_dims = 0 : i32}> : (tensor<3xi32>, tensor<2xi64>) -> tensor<2xi32>
// CHECK-NEXT:      %18 = "tfl.pseudo_const"() <{value = dense<1> : tensor<1xi64>}> : () -> tensor<1xi64>
// CHECK-NEXT:      %19 = "tfl.gather"(%15, %18) <{axis = 0 : i32, batch_dims = 0 : i32}> : (tensor<3xi32>, tensor<1xi64>) -> tensor<1xi32>

  %6 = "tfl.concatenation"(%1, %0) {axis = -1 : i32, fused_activation_function = "NONE"} : (tensor<1x1xf32>, tensor<1x1xf32>) -> tensor<1x2xf32>
  %7 = "quantfork.stats"(%6) {layerStats = dense<[-0.440728068, 0.26483655]> : tensor<2xf32>} : (tensor<1x2xf32>) -> tensor<1x2xf32>

  fbb.String("string_separator", string_separator_str);

  auto axis = mlir::dyn_cast_or_null<IntegerAttr>(attrs.get("axis"));
  if (!axis) {
    return func.emitError() << "'axis' attribute is not set or not an integer";
  }
  fbb.Int("axis", axis.getInt());

  auto reduction_type =
      mlir::dyn_cast_or_null<StringAttr>(attrs.get("reduction_type"));