%1 = "tfl.quantize"(%arg0) {qtype = tensor<1x224x224x3x!quant.uniform<u8:f32, 7.812500e-03:128>>} : (tensor<1x224x224x3xf32>) -> tensor<1x224x224x3x!quant.uniform<u8:f32, 7.812500e-03:128>>

  func.return %24 : tensor<1x4xf32>

  func.return %24 : tensor<1x4xf32>

  for (auto arg_num : input_mapping) {
    const Tensor* tensor;
    if (auto it = variable_snapshots.find(arg_num);
        it != variable_snapshots.end()) {
      tensor = it->second;
    } else {
      tensor = inputs[arg_num - num_missing_prefix_ctx_inputs];
    }

    // The input tensor can have the following cases.
    // 1. Tensor with PjRtTensorBuffer, containing a PjRtBuffer. This case

func.func @main(tensor<4x4x4xf32>, tensor<4x4xf32>, tensor<4x4xf32>, tensor<4x4xf32>, tensor<4x4xf32>, tensor<4x4xf32>, tensor<4x4xf32>, tensor<4x4xf32>, tensor<4x4xf32>, tensor<4xf32>, tensor<4xf32>, tensor<4xf32>, tensor<4xf32>, tensor<4xf32>, tensor<4xf32>, tensor<4xf32>, tensor<4x4xf32>, tensor<4xf32>, tensor<4x4xf32>, tensor<4x4xf32>, tensor<4x4xf32>, tensor<4x4xf32>) -> tensor<4x4x4xf32> {
// CHECK: {
// CHECK-NEXT:   version: 3,

  // quantization aware training or calibration, for the remaining tensors.
  std::vector<std::pair<std::optional<double>, std::optional<double>>>
      input_ranges;

  // Whether to disable setting the quantization parameters of the input nodes
  // using input ranges.
  bool disable_set_input_nodes_quantization_params = false;

  // The default ranges can be used when a tensor doesn't have quantization

 ...

    status->status = tensorflow::errors::InvalidArgument(
        "DLPack doesn't support ", handle->TypeString(), " tensor");
    return nullptr;
  }
  const tensorflow::Tensor* tensor;
  status->status = handle->Tensor(&tensor);
  if (!status->status.ok()) {
    return nullptr;
  }
  return tensor;
}

// Deleter for DLManagedTensor
void DLManagedTensorDeleter(DLManagedTensor* arg) {
  TfDlManagedTensorCtx* owner =

  std::vector<const Tensor*> inputs(input_arg_size);
  std::deque<Tensor> inputs_storage;
  std::vector<VariableInfo> variable_infos;
  int offset = input_arg_shape_and_dtype.size();

  for (int i = 0; i < input_handles.size(); i++) {
    const TensorHandle* th = input_handles[i];
    const Tensor* t;
    // Handle owns the tensor.
    TF_RETURN_IF_ERROR(th->Tensor(&t));

      %replicated_output:2 = "tf.TPUReplicatedOutput(%identity) : (tensor<i32>) -> (tensor<i32>, tensor<i32>)
      return %replicated_output#0, %replicated_output#1 : tensor<i32>, tensor<i32>
    }
    ```

    will be transformed into:

    ```mlir
    func @tpu_computation(%arg0: tensor<i32>, %arg1: tensor<i32>) -> (tensor<i32>, tensor<i32>) {