// time-majored. shrink_axis_mask: should be 0b010 = 2 or 0b001 = 1 if it's
  // time-majored.
  SmallVector<int64_t, 2> last_output_shape({batch, n_output});

  SmallVector<int32_t, 3> end({0, 0, 0});
  SmallVector<int32_t, 3> strides({1, 1, 1});
  SmallVector<int32_t, 3> begin;

  int64_t new_axis_mask = 0;
  int64_t ellipsis_mask = 0;
  int64_t begin_mask;
  int64_t end_mask;

            .getValue(),
        ftype);

    func->setAttrs(graph_func->getAttrs());

    llvm::SmallVector<Type> arg_types;
    llvm::SmallVector<Type> res_types;
    llvm::SmallVector<DictionaryAttr> arg_attrs;
    llvm::SmallVector<DictionaryAttr> res_attrs;
    if (failed(FilterTfgSpecificArgResultAttributes(
            getContext(), ftype.getInputs(), graph_func.getAllArgAttrs(),

  const UniformQuantizedPerAxisType quantized_type =
      CreateI8F32UniformQuantizedPerAxisType(
          UnknownLoc::get(&ctx_), ctx_,
          /*scales=*/SmallVector<double, 2>{1.0, 1.0},
          /*zero_points=*/SmallVector<int64_t, 2>{0, 0},
          /*quantization_dimension=*/0);

  // Storage type of `i8` is currently verifiable as `unsigned` in `Types.cpp`.

}

// Return a vector of the return indices.
llvm::SmallVector<llvm::SmallVector<int, 4>> GetReturnIndicesVec(
    const llvm::SmallVector<TF::IfRegionOp, 8>& if_op_segment) {
  llvm::SmallVector<llvm::SmallVector<int, 4>> return_indices_vec;
  for (auto it = if_op_segment.begin(); it != if_op_segment.end(); ++it) {
    llvm::SmallVector<int, 4> indices_to_keep_vec =
        GetReturnIndicesToKeep(*it, if_op_segment);

func::FuncOp createLstmCompositeFunc(mlir::Builder* builder, bool ln,
                                     bool cifg) {
  SmallVector<int64_t, 2> input_shape{1, 2};
  SmallVector<int64_t, 2> weight_shape{3, 12};
  SmallVector<int64_t, 1> bias_shape{2};
  SmallVector<int64_t, 2> projection_shape{1, 2};
  SmallVector<int64_t, 1> layer_norm_scale{4};
  SmallVector<int64_t, 2> output_shape{1, 2};
  auto input_type = RankedTensorType::get(input_shape, builder->getF32Type());

  assert(devices.find(tensorflow::GetDeviceAliasForLogicalCore(0))
             ->getSecond()
             .size() == num_replicas);

  llvm::SmallVector<std::pair<ValueRange, Type>, 8> new_replicated_inputs;
  llvm::SmallVector<Value, 8> new_packed_inputs;
  llvm::SmallVector<llvm::SmallVector<Value, 8>, 8> replicated_inputs;
  replicated_inputs.reserve(replicate.GetNumReplicatedBlockArguments());

AnonymousIteratorV3Op CreateIterator(OpBuilder builder,
                                     llvm::ArrayRef<Type> dataset_types,
                                     ReduceDatasetOp reduce_dataset) {
  llvm::SmallVector<Attribute, 2> shape_attrs;
  llvm::SmallVector<Attribute, 2> type_attrs;
  for (Type type : dataset_types) {
    shape_attrs.push_back(
        TF::ShapeAttr::get(builder.getContext(), mlir::cast<ShapedType>(type)));

};

// Compute the permutation that would take `arr` to the identity.
llvm::SmallVector<int32_t> InvertPermutation(llvm::SmallVector<int32_t> arr) {
  llvm::SmallVector<int32_t> inverse_arr(arr.size());
  for (int32_t i = 0; i < arr.size(); ++i) {
    inverse_arr[arr[i]] = i;
  }
  return inverse_arr;
}

llvm::SmallVector<int64_t> PermuteShape(llvm::ArrayRef<int64_t> shape,

  // TFL::UnsortedSegmentProdOp.
  llvm::SmallVector<int32_t, 4> flattened_out_segids =
      llvm::SmallVector<int32_t, 4>(operand_rank, static_cast<int32_t>(-1));
  for (int64_t i : dot_dimensions_info.out_dimensions().AxesArray()) {
    flattened_out_segids[i] = 0;
  }
  llvm::SmallVector<int32_t, 4> flattened_contracting_segids =
      llvm::SmallVector<int32_t, 4>(operand_rank, static_cast<int32_t>(-1));

    StringAttr conv_padding, Value &padding, int num_dims) {
  Value zero_rank1 = CreateConstValue<int32_t>(builder, loc, {1}, {0});
  SmallVector<Value> temp_padding_values{zero_rank1, zero_rank1};

  auto reshape_op = [&](Value value, const SmallVector<int64_t> &shape) {
    const int64_t rank = shape.size();
    return builder.create<TF::ReshapeOp>(