MutableArrayRef<int64_t> input_shape, ArrayRef<int64_t> sparse_begin,
    ArrayRef<int64_t> sparse_end, ArrayRef<int64_t> sparse_strides,
    int32_t begin_mask, int32_t end_mask, int32_t ellipsis_mask,
    int32_t new_axis_mask, int32_t shrink_axis_mask,
    SmallVectorImpl<int64_t> *begin, SmallVectorImpl<int64_t> *end,
    SmallVectorImpl<int64_t> *stride) {
  int64_t num_sparse_indices = sparse_begin.size();

// pointer.
void CallFunctionsBase(::testing::benchmark::State& state) {
  int64_t sum = 0;
  typedef int64_t (*Func)(int64_t a, int64_t b);
  volatile Func f_raw = AddRaw;
  Func f = f_raw;
  size_t i = 0;
  for (auto dummy : state) {
    sum = f(sum, i);
    i++;
  }
  // volatile int64_t result = sum;
}

BENCHMARK(CallFunctions)->Arg(1 << 10);
BENCHMARK(CallFunctionsIndirect)->Arg(1 << 10);

    auto output = op->getResult(0);
    auto output_type =
        mlir::dyn_cast_or_null<mlir::RankedTensorType>(output.getType());
    if (!output_type || !output_type.hasStaticShape()) return false;

    *count = output_type.getNumElements();
    return true;
  }

  static bool GetInputTensorTotalSize(mlir::Operation* op, int64_t* count) {
    int64_t total_count = 0;

  const int64_t max_us = (options.max_micros <= 0 && options.max_iters <= 0)
                             ? Options::kDefaultMicros
                             : options.max_micros;
  // NOLINTNEXTLINE
  printf("Running benchmark for %lld us\n", static_cast<long long>(max_us));
  const int64_t start_us = NowMicros();
  int64_t iters = 0;
  while (true) {

      /*shape=*/Create1DConstValue<int64_t>(builder, loc, {operand_rank}),
      /*value=*/CreateScalarConstValue<int64_t>(builder, loc, 0));

  // Converts start_index_map proto to tensor.
  const int64_t index_map_size = dimension_numbers.start_index_map().size();
  SmallVector<int64_t> indices(index_map_size);
  for (int64_t i = 0; i < index_map_size; i++) {

namespace tensorflow {
namespace tfrt_compiler {
namespace {

constexpr int64_t kDefaultCheapCost = 1;

int64_t GetRankedTensorSize(mlir::TensorType type) {
  auto shape = type.getShape();

  int64_t size = 1;
  for (int64_t dim : shape) {
    // For unknown dimensions, use 1 as the size because it is usually the batch
    // dimension.
    //

  llvm::SmallVector<std::pair<int64_t, llvm::SmallVector<string, 4>>, 4>
      arg_grads;
  llvm::SmallVector<std::pair<int64_t, int64_t>, 4> ret_forward_input;
};

// Updates a called function's input signature by adjusting resource types, and
// adding required gradient arguments.
void ChangeFunctionInputSignature(
    func::FuncOp func,
    const llvm::SmallDenseMap<int64_t, llvm::SmallVector<string, 4>>& grads,

inline ShapedType GetTransposedType(Value input,
                                    llvm::ArrayRef<int64_t> permutation_array) {
  auto input_type = input.getType().cast<ShapedType>();
  if (permutation_array.size() != input_type.getRank()) {
    return nullptr;
  }
  llvm::SmallVector<int64_t> transposed_shape(permutation_array.size());
  for (int64_t i = 0; i < permutation_array.size(); ++i) {

  const ArrayRef<int64_t> lhs_contracting_dims =
      dot_dimension_numbers.getLhsContractingDimensions();
  const ArrayRef<int64_t> rhs_contracting_dims =
      dot_dimension_numbers.getRhsContractingDimensions();
  const int64_t input_rank =
      mlir::dyn_cast<ShapedType>(dot_general_op.getOperand(0).getType())
          .getRank();
  const int64_t filter_rank =

  // Figure out output shapes.
  SmallVector<int64_t, 4> concat_out_shape;
  SmallVector<int64_t, 4> pack_out_shape;

  const int64_t rank = input_type.getRank();
  int64_t pack_axis = pack_op.getAxis();
  size_t count = pack_inputs.size();
  if (pack_axis < 0) {
    pack_axis += rank;
  }

  // Concat out shape.
  for (int i = 0; i < rank; ++i) {
    int64_t dim_size = input_type.getDimSize(i);