xs.push_back(Placeholder(scope, DT_DOUBLE, Placeholder::Shape(shape)));
  xs.push_back(Placeholder(scope, DT_DOUBLE, Placeholder::Shape(shape)));
  auto tmp = Stack(scope, xs, Stack::Axis(0));
  auto y = Unstack(scope, tmp, 2, Unstack::Axis(0));
  double max_error;
  TF_ASSERT_OK((ComputeGradientError<double, double, double>(
      scope, xs, {shape, shape}, y.output, {shape, shape}, &max_error)));
  EXPECT_LT(max_error, 1e-10);

  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;
}

//
// Description:
//   Given a tensor `input`, this operation inserts a dimension of 1 at the
//   dimension index `axis` of `input`'s shape. The dimension index `axis`
//   starts at zero; if you specify a negative number for `axis` it is counted
//   backward from the end.
//
//   This operation is useful if you want to add a batch dimension to a single

// 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.