return absl::InternalError(
          "DepthwiseConv2D doesn't get any benefit of quantization in XLA.");
    } else if (function_name.contains("conv2d")) {
      // For Conv2D, the channel dimension must be static to calculate the
      // feature group count.
      if (!HasStaticShapeAtDims(call_op->getOperand(0), /*dims=*/3)) {
        return absl::InternalError(

  if (av_tensor == nullptr || av_tensor->GetBuffer() == nullptr) {
    return absl::InternalError("Input tensor does not have PjRtBuffer.");
  }
  auto* c_api_buffer =
      dynamic_cast<xla::PjRtCApiBuffer*>(av_tensor->GetBuffer().get());
  if (c_api_buffer == nullptr) {
    return absl::InternalError(
        "The PjRtBuffer in the tensor is not type PjRtCApiBuffer.");
  }
  return c_api_buffer->c_buffer();
}

                                          absl::MakeSpan(exported_names),
                                          context, import_options, &bundle);
  if (!module.status().ok()) {
    return absl::InternalError(absl::StrCat("Failed to import SavedModel: ",
                                            module.status().message()));
  }

  mlir::OwningOpRef<mlir::ModuleOp> module_ref = std::move(module).value();

namespace stablehlo::quantization::io {

absl::StatusOr<std::string> GetLocalTmpFileName(tsl::Env* const env) {
  std::string tmp_fname{};
  if (!env->LocalTempFilename(&tmp_fname)) {
    return absl::InternalError("Failed to create tmp file name.");
  }

  return tmp_fname;
}

absl::StatusOr<std::string> GetLocalTmpFileName() {
  return GetLocalTmpFileName(tsl::Env::Default());
}

                            tpu::TPUCompileMetadataProto& metadata_proto) {
  auto main_fn = module.lookupSymbol<mlir::func::FuncOp>(kEntryFuncName);
  if (!main_fn) {
    return absl::InternalError("Could not find main function in MLIR Module.");
  }

  mlir::FunctionType func_type = main_fn.getFunctionType();
  for (auto input_type : func_type.getInputs()) {
    tensorflow::TensorShape tensor_shape;

    return absl::InternalError("Couldn't import flatbuffer to MLIR.");
  }

  PassManager pm((*module)->getName(), OpPassManager::Nesting::Implicit);
  pm.addPass(TFL::CreateDenseToSparsePass());

  if (failed(pm.run(module.get()))) {
    LOG(ERROR) << "Failed to sparsify: "
               << statusHandler.ConsumeStatus().message();
    return absl::InternalError(absl::StrCat(

// Creates a temporary directory on an environment defined by the implementation
// of `tsl::Env` and returns its path. Returns an InternalError status if
// failed.
absl::StatusOr<std::string> CreateTmpDir(tsl::Env* env);

// Creates a temporary directory and returns its path. Returns an InternalError
// status if failed. The file system used will be the default environment
// returned by `tsl::Env::Default`.

    exported_names.push_back(key);
  }

  if (failed(mlir::tf_saved_model::FreezeVariables(
          module_op, saved_model_bundle->GetSession()))) {
    return absl::InternalError("Failed to freeze variables.");
  }

  // Run legalize TF to StableHLO pass to convert `tf.Const` and
  // `tf.Const`->`tf.Cast` patterns after variable freezing. The TF shape

    mlir::StringRef id = aggregator_op.getId();
    auto iter = statistics_map.find(id);
    if (iter == statistics_map.end()) {
      status = absl::InternalError(
          absl::StrFormat("Calibrated data does not exist. Cannot find "
                          "statistics. value for id: %s",
                          id));
      return;
    }

    const auto status = quantization::PreprocessAndFreezeGraph(
        module, module.getContext(), session);
    if (!status.ok()) {
      return status_handler.Combine(
          absl::InternalError("Failed to preprocess & freeze TF graph."));
    }

    // TODO: b/264218457 - Refactor the component below once StableHLO Quantizer
    // can run DRQ. Temporarily using TF Quantization for StableHLO DRQ.