int64_t num_updates = indices_type.getDimSize(0);
      // For TF::TensorScatterUpdateOp, `indices` must have at least 2 axes:
      // `(num_updates, index_depth)`. Reshape indices and updates if necessary.
      if (std::is_same<TfOp, TF::TensorScatterUpdateOp>::value &&
          indices_type.getRank() == 1 && updates_type.getRank() == 1 &&
          index_depth == 1 && num_updates == 1) {

    indices[i] = dimension_numbers.start_index_map()[i];
  }

  // Fill elements from start_indices with start_index_map
  Value scattered_start_indices = builder.create<TF::TensorScatterUpdateOp>(
      loc, empty_start_indices,
      /*indices=*/
      builder.create<TF::ReshapeOp>(
          loc, RankedTensorType::get({index_map_size, 1}, builder.getI64Type()),

    TypeID::get<TF::StopGradientOp>(),
    TypeID::get<TF::StridedSliceOp>(),
    TypeID::get<TF::StridedSliceGradOp>(),
    TypeID::get<TF::SumOp>(),
    TypeID::get<TF::TanhGradOp>(),
    TypeID::get<TF::TensorScatterUpdateOp>(),
    TypeID::get<TF::TileOp>(),
    TypeID::get<TF::TopKV2Op>(),
    TypeID::get<TF::_UnaryOpsCompositionOp>(),
    TypeID::get<TF::UnsortedSegmentMaxOp>(),
    TypeID::get<TF::UnsortedSegmentMinOp>(),

  auto indices = builder.create<TF::ReshapeOp>(scatter.getLoc(),
                                               scatter.getIndices(), shape);
  Value tensor_scatter_update = builder.create<TF::TensorScatterUpdateOp>(
      scatter.getLoc(), buffer, indices, scatter.getTensor());
  scatter.getOutputHandle().replaceAllUsesWith(tensor_scatter_update);
  scatter.erase();
  auto size = it->getSecond();