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

  %axis = "tf.Const"() {device = "", value = dense<1> : tensor<i32>} : () -> tensor<i32>
  // tf.ConcatV2 accepts a variadic operand. The last operand should be compile
  // time constant.
  %0 = "tf.ConcatV2"(%arg0, %arg0, %arg0, %arg0, %axis) : (tensor<16x1xf32>, tensor<16x1xf32>, tensor<16x1xf32>, tensor<16x1xf32>, tensor<i32>) -> tensor<16x4xf32>

  // CHECK-SAME: (tensor<i32>, tensor<i32>, tensor<i32>) -> tensor<2xi32>
  %axis = "tf.Const"() {value = dense<0> : tensor<i32>} : () -> tensor<i32>
  %0 = "tf.VariableV2"() {container = "", shape = #tf_type.shape<>, shared_name = "x"} : () -> tensor<!tf_type.int32ref>
  %1 = "tf.ConcatV2"(%0, %0, %axis) : (tensor<!tf_type.int32ref>, tensor<!tf_type.int32ref>, tensor<i32>) -> tensor<2xi32>
  func.return %1 : tensor<2xi32>

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