// expected-error @+1 {{op attribute 'axis' should be in range [-rank - 1, rank + 1), got rank = 1, and axis = 3}}
  %0 = "tfl.pack"(%arg0, %arg1) {axis = 3 : i32, values_count = 2 : i32} : (tensor<2xi32>, tensor<2xi32>) -> tensor<2x2xi32>
  func.return %0 : tensor<2x2xi32>
}

// -----

func.func @unpack(%arg0: tensor<2x3xi32>) -> tensor<2xi32> {
  // CHECK: "tfl.unpack"(%arg0) <{axis = 1 : i32, num = 3 : i32}>

  if (!axes_attr) return false;

  for (auto a : axes_attr.getValues<APInt>()) {
    int64_t axis = a.getSExtValue();
    if (axis < 0) {
      axis += type_shape.size();
    }
    if (axis < 0 || axis >= type_shape.size()) {
      // `axis` is not a valid axis in input.
      return false;
    }
    if (type_shape[axis] != 1) {
      return false;
    }
  }

  return true;
}

// to expand at the given `axis`.
Type InferExpandDimsType(Type ty, int64_t axis, Builder *builder) {
  auto ranked_ty = mlir::dyn_cast<RankedTensorType>(ty);

  // Unranked type.
  if (!ranked_ty) return ty;

  auto shape = llvm::to_vector<4>(ranked_ty.getShape());
  if (axis < 0) axis += ranked_ty.getRank() + 1;

  shape.insert(shape.begin() + axis, 1);

  auto values = tf_concat_op.getValues();
  auto output_type = tf_concat_op.getOutput().getType();
  // Extract axis attribute from constant axis tensor
  ElementsAttr axis;
  if (!matchPattern(tf_concat_op.getAxis(), m_Constant(&axis)))
    return failure();
  IntegerAttr axis_int = ExtractSingleElementAsInteger(axis);

  // "axis" operand could be a i64 tensor. Resolve it here.
  IntegerAttr axis_i32;

      // $group_assignment
      CompileTimeConstantOperand<TF::CrossReplicaSumOp, 1>,
      CompileTimeConstantOperand<TF::CumprodOp, 1>,              // $axis
      CompileTimeConstantOperand<TF::CumsumOp, 1>,               // $axis
      CompileTimeConstantOperand<TF::CumulativeLogsumexpOp, 1>,  // $axis
      // $filter_sizes
      CompileTimeConstantOperand<TF::DepthwiseConv2dNativeBackpropFilterOp, 1>,
      // $input_sizes

  [(IsNotInLiftedFunc $res), (IsConstTensor $b)], [], (addBenefit 1)>;

def LiftGather : Pat<
  (TF_GatherV2Op:$res $params, $indices, $axis, $batch_dims),
  (LiftAsTFPartitionedCall<"composite_gather_fn">
    (ArgumentList $params, $indices, $axis),
    (ResultList $res),
    (NamedAttributeList
      (NamedAttr<"batch_dims"> $batch_dims))),
  [(IsNotInLiftedFunc $res), (IsConstTensor $params)], [], (addBenefit 1)>;

  // CHECK: %[[MATMUL_PACKED:.*]] = "tf.Pack"(%[[MATMUL_1]], %[[MATMUL_2]], %[[MATMUL_3]], %[[MATMUL_4]], %[[MATMUL_5]], %[[MATMUL_6]]) <{axis = 0 : i64}> : (tensor<4x6xf32>, tensor<4x6xf32>, tensor<4x6xf32>, tensor<4x6xf32>, tensor<4x6xf32>, tensor<4x6xf32>) -> tensor<6x4x6xf32>

    int64_t depth = depth_attr.getValues<APInt>()[0].getSExtValue();
    int64_t axis = op.getAxis();
    if (axis == -1) axis = indices_shape.size();

    llvm::SmallVector<int64_t, 4> broadcast_dims(indices_shape.size());
    std::iota(broadcast_dims.begin(), broadcast_dims.begin() + axis, 0);
    std::iota(broadcast_dims.begin() + axis, broadcast_dims.end(), axis + 1);

    llvm::SmallVector<int64_t, 4> output_dims =

// Concat op patterns.
//===----------------------------------------------------------------------===//

def ConvertToConcatV2 : Pat<
  (TF_ConcatOp:$src $axis, $inputs), (TF_ConcatV2Op:$dest $inputs, $axis),
  [], [(CopyAttrs $src, $dest)]>;

//===----------------------------------------------------------------------===//
// Div op patterns.

    }
    // Make sure that the axis in `expand_op` is constant.
    if (auto const_op =
            llvm::dyn_cast<TF::ConstOp>(expand_op.getDim().getDefiningOp())) {
      expand_axis = (*mlir::cast<DenseElementsAttr>(const_op.getValue())
                          .getValues<APInt>()
                          .begin())
                        .getSExtValue();
      // Canonicalize axis. Some TF python functions, such as