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,

// required to collect tensor statistics.
void AddPreCalibrationPasses(
    OpPassManager& pm,
    const ::stablehlo::quantization::CalibrationOptions& calibration_options,
    const ::stablehlo::quantization::QuantizationSpecs& specs,
    const ::stablehlo::quantization::DebuggerConfig& debugger_config);

// Adds passes for static-range quantization post-calibration. Utilizes tensor

  // functionality will be supported in the future.

  // Create a host tensor and copy data into it.
  //
  // Raises an error if shape or dtype are incompatible with T.
  // TODO(b/189458441): Update this when we decide on the representation of
  // shape and dtype in this API.
  // Disclaimer: This API is subject to change as we add support for creating
  // device tensors b/187222691 and enable buffer re-use b/187223179.

opt<std::string> input_shapes(
    "tf-input-shapes",
    llvm::cl::desc(
        "Input tensor shapes. Shapes for different tensors are separated by "
        "':', and dimension sizes for the same tensor are separated by ','"),
    llvm::cl::init(""));

// NOLINTNEXTLINE
opt<std::string> output_arrays(
    "tf-output-arrays", llvm::cl::desc("Output tensor names, separated by ','"),
    llvm::cl::init(""));

// NOLINTNEXTLINE

    base_offset *= dimension;
  }

  return contiguous_offset;
}

// Performs transposition of a tensor represented as a contiguous element array.
// Assumes row-major order. The shape of the input tensor and the desired
// permutation is registered during construction, and calling `TransposeValues`
// returns the transposed tensor values.
class DenseElementsTransposer {
 public:

  // 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

// initialized, the value will not be present.
using ResourceVarsSnapshot = absl::flat_hash_map<int, std::optional<Tensor>>;

// Takes a snapshot of the values of resource variable arguments, whose indices
// are specified in `variable_indices` argument. We snapshot tensors that back
// resource variables since concurrent updates may modify the shape, and it is

  // not applied regardless of the op support. Currently, it is supported for
  // XLA opset for SRQ on weight tensors (not activation),
  // and Uniform Quantized opset .
  optional bool enable_per_channel_quantization = 10;

  // Enables two inputs of an operation to be both tensors.
  // Currently supports MatMul and BatchMatMul ops for XLA.
  // TODO(b/263528090): Check the condition when this feature is beneficial.

      }
      return %graph#0, %graph#1, %graph#2, %graph#3 : !tf_res, !tf_res, tensor<f32>, tensor<f32>
    }
    ```

    After:

    ```mlir
    func @while_body(%arg0: !tf_res, %arg1: !tf_res, %arg2: tensor<f32>, %arg3: tensor<f32>, %chain_0: tensor<i32>, %chain_1: tensor<i32>) -> (!tf_res, !tf_res, tensor<f32>, tensor<f32>, tensor<i32>, tensor<i32>) {