%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 =

  }
  return tpu_shape.AsCpp<xla::Shape>();
}

// Given a tensor, returns the shape of its representation on device,
// fully padded. Contents of `shape` are undefined on error.
Status TpuPaddedShapeFn(const Tensor& tensor, xla::Shape* shape) {
  const tensorflow::XlaTensor* xla_tensor =
      tensorflow::XlaTensor::FromTensor(&tensor);
  if (xla_tensor == nullptr) {
    return errors::InvalidArgument(

  for (int component_index = 0; component_index < tensors_.size();
       ++component_index) {
    // TODO(allenl): Add a C API for summarizing tensors. Currently custom
    // devices limiting themselves to a C API (for ABI compatibility) would need
    // to implement summarization for component tensors themselves.
    ImmediateExecutionTensorHandle* component =
        tensorflow::unwrap(tensors_[component_index].get());