llvm::SmallVector<std::pair<int64_t, llvm::SmallVector<Value, 4>>, 4>
AnnotateCompileOpAndGetExecuteArgToWhileArgsMapping(
    TF::WhileRegionOp while_op, tf_device::ReplicateOp replicate,
    TF::TPUExecuteAndUpdateVariablesOp execute,
    tf_device::LaunchOp compile_launch) {
  Region& body = while_op.getBody();
  Region& cond = while_op.getCond();

  llvm::SmallVector<std::pair<int64_t, llvm::SmallVector<Value, 4>>, 4> mapping;

std::vector<DataType> DataTypeSetToVector(DataTypeSet set);

// Returns a vector of shapes intended to be "interesting" test cases.
// Currently, this returns scalar, 1D vector, 2D matrix, and a 4D tensor shapes
std::vector<std::vector<int64_t>> InterestingShapes();

// Returns a TensorHandle of `dtype` and `shape`, filled with `value`.
// `dtype` must be an integer dtype, float, or double.
// If a TensorHandle cannot be created successfully, this function will

}

template <>
inline std::vector<int64_t> GetVector(DenseElementsAttr elements) {
  auto type = elements.getType();
  auto elemType = type.getElementType();
  if (elemType.isSignlessInteger(64)) {
    auto vec = llvm::to_vector(llvm::map_range(
        elements.getValues<APInt>(),
        [&](APInt value) -> int64_t { return value.getSExtValue(); }));
    return std::vector<int64_t>(vec.begin(), vec.end());
  }

      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`.
  EXPECT_TRUE(quantized_type.getStorageType().isSignlessInteger(8));
}

static mlir::ShapedType ReversePermuteShapedType(
    mlir::ShapedType type, ArrayRef<int64_t> permutation) {
  if (!type.hasRank()) return type;

  auto shape = type.getShape();
  SmallVector<int64_t, 4> new_shape(shape.size());

  for (int i = 0; i < permutation.size(); ++i) {
    int64_t index = permutation[i];
    assert(index < shape.size());
    new_shape[index] = shape[i];
  }

LogicalResult CreateInitBufferValue(ArrayRef<int64_t> element_shape,
                                    int64_t max_size, Operation* op,
                                    Type element_dtype, OpBuilder builder,
                                    Value* buffer);

// Same as above, but uses a Value as max_size and check if it is a constant.
LogicalResult CreateInitBufferValue(ArrayRef<int64_t> element_shape,

  static std::pair<int64_t, int64_t> GetBiasDimAndSliceSize(
      ArrayRef<int64_t> filter_shape, AffineOpType op) {
    // Channel dimension index is specified as op property
    auto channel_index_iter = filter_shape.begin();
    std::advance(channel_index_iter, op.GetChannelDimIndex());
    // The slide size is the size of the data in higher dimensions.
    int64_t slice_size =

  // Get a list of the names of functions that have been registered.
  virtual std::vector<string> ListFunctionNames() = 0;

  struct CacheStats {
    int64_t kernel_cache_size;
    int64_t device_cache_size;
    std::map<std::string, int64_t> func_kernel_cache_entries;
    int64_t local_rendezvous_cache_active_size;
  };
  virtual CacheStats GetCacheStats() = 0;

    flatbuffers::FlatBufferBuilder* builder, const tflite::Model* input_model,
    const tflite::TensorType& inference_type,
    const absl::flat_hash_set<std::string>& denylisted_ops,
    const CustomOpMap& custom_op_map,
    int64_t minimum_elements_for_weights = 1024,
    bool disable_per_channel = false, bool weight_only_quantization = false,
    bool legacy_float_scale = false);

// Overloading methods to support old quantizer versions API

namespace mlir {
namespace TF {

llvm::SmallVector<ResourceHandleValueAndId, 4>
_TfrtGetResourceOp::GetResourceHandleValueAndIdList(
    llvm::SmallDenseMap<ResourceHandle, int64_t> &resource_handle_id_map,
    int64_t &next_id) {
  llvm::SmallVector<ResourceHandleValueAndId, 4> resource_vec;
  llvm::StringRef device = GetDeviceOrEmpty(getOperation());

  for (const auto &iter : llvm::enumerate(getResults())) {