llvm::SmallVector<int64_t>({num_updates, index_depth}),
            rewriter.getI32Type());
        new_indices = rewriter.create<TF::ReshapeOp>(
            loc,
            RankedTensorType::get({num_updates, index_depth},
                                  indices_type.getElementType()),
            indices, indices_shape);
        auto updates_shape =
            BuildIntArrayConstOp(builder, rewriter,

TF::ConstOp Create1DStringConst(const ArrayRef<std::string> str_values,
                                const Location loc, OpBuilder& builder) {
  const auto tensor_type =
      RankedTensorType::get(/*shape=*/{static_cast<int64_t>(str_values.size())},
                            /*elementType=*/builder.getType<TF::StringType>());

  return builder.create<TF::ConstOp>(
      loc, DenseStringElementsAttr::get(

                                                     dense_data.data(), inst,
                                                     &builder, &s_param);
          auto compressed_data_type = RankedTensorType::get(
              {static_cast<int64_t>(compressed_data.size())},
              builder.getF32Type());
          auto new_value = DenseElementsAttr::get<float>(compressed_data_type,

    return custom_call.getCallTargetName() == "mhlo.erf";
  }
  return llvm::isa<chlo::ErfOp>(op);
}

// Builds a reference implementation of non-approximate GELU.
func::FuncOp BuildGELUDecomposition(RankedTensorType type,
                                    PatternRewriter& rewriter,
                                    Block* insertion_point) {
  rewriter.setInsertionPointToStart(insertion_point);

                                  mlir::OpBuilder* builder,
                                  mlir::TF::SplitOp* split_op) {
  // Creates a const op to hold split dimension value.
  auto split_dim_type =
      mlir::RankedTensorType::get({}, builder->getIntegerType(32));
  auto split_dimension_attr =
      mlir::DenseElementsAttr::get(split_dim_type, split_dimension);
  auto split_dimension_op = builder->create<mlir::TF::ConstOp>(

  ElementsAttr layer_stats = DenseFPElementsAttr::get(
      RankedTensorType::get({2}, b.getF32Type()), {min_maxs[0], min_maxs[1]});
  ElementsAttr axis_stats;
  IntegerAttr axis;

  if (info.params_size() > 1) {
    SmallVector<int64_t, 4> axis_stats_shape{info.params_size(), 2};
    axis_stats = DenseFPElementsAttr::get(
        RankedTensorType::get(axis_stats_shape, b.getF32Type()), min_maxs);

def GetHLOAxisFromTFAxis : NativeCodeCall<
  "GetHLOAxisFromTFAxis("
  "$0, $1.getType().cast<RankedTensorType>().getRank(), &$_builder)">;

// Same as the above but with $1 of type operand_range from variadic TensorFlow
// input.
def GetHLOAxisFromTFAxisVariadic : NativeCodeCall<
  "GetHLOAxisFromTFAxis("
  "$0, (*$1.begin()).getType().cast<RankedTensorType>().getRank(), "
  "&$_builder)">;

def CastElementsToI64Elements : NativeCodeCall<

};

LogicalResult CheckFusableLayerNormalizedLstmCellSimple(
    func::FuncOp lstm_func) {
  for (int i = 0; i < 5; ++i) {
    auto input = lstm_func.getArgument(i);
    auto input_type = mlir::dyn_cast_or_null<RankedTensorType>(input.getType());
    if (!input_type) {
      lstm_func.emitWarning(
          "we cannot fuse this lstm func because all the inputs have not "
          "ranked tensor type.");
      return failure();
    }

  if constexpr (std::is_integral_v<T>) {
    auto shape_type =
        RankedTensorType::get(shape, builder.getIntegerType(sizeof(T) * 8));

    const auto attr = DenseIntElementsAttr::get(shape_type, values);
    return builder.create<TF::ConstOp>(loc, attr);
  }

  const auto type = RankedTensorType::get(shape, builder.getF32Type());
  const auto value_attr = DenseFPElementsAttr::get(type, values);

    if (cur_rank != 4 || cur_shape[2] == 1) return failure();
    TensorType new_shape = RankedTensorType::get(
        {cur_shape[0], cur_shape[1], 1, cur_shape[2] * cur_shape[3]},
        attr.getElementType());

    // Inserts a reshape op.
    auto shape_spec_type =
        RankedTensorType::get({cur_rank}, rewriter.getIntegerType(64));
    auto new_shape_const_attr =