tags
    ).meta_info_def.function_aliases

  def _create_matmul_model(
      self,
      input_shape: Sequence[int],
      weight_shape: Sequence[int],
      saved_model_path: str,
      bias_fn: Optional[ops.Operation] = None,
      activation_fn: Optional[ops.Operation] = None,
  ) -> module.Module:
    class MatmulModel(module.Module):
      """A simple model with a single matmul.

      function_aliases = GetFunctionAliases(saved_model_path, tags);
  if (!function_aliases.ok()) {
    return absl::InternalError(absl::StrCat(
        "Failed to get function alias: ", function_aliases.status().message()));
  }

  absl::StatusOr<mlir::OwningOpRef<mlir::ModuleOp>> module =
      ImportAndPreprocessSavedModel(
          saved_model_path, signature_keys, tags, context.get(),

        func=functools.partial(
            _run_calibration,
            saved_model_path,
            signature_keys,
            tags,
            force_graph_mode_calibration,
            dataset_file_map,
        ),
        error_msg=(
            f'Failed to run calibration on model "{saved_model_path}",'
            f' signature_keys: {signature_keys}, tags: {tags}.'
        ),
    )

_DYNAMIC_RANGE_DEFAULT_MIN_NUM_ELEMENTS_FOR_WEIGHTS = 1024


def _is_qat_saved_model(saved_model_path: str):
  """Checks if the SavedModel is QAT-enabled by looking for 'FakeQuant' ops."""
  saved_model_proto = saved_model_loader.parse_saved_model(saved_model_path)
  for meta_graph in saved_model_proto.meta_graphs:
    if any(
        node.op.startswith('FakeQuant') for node in meta_graph.graph_def.node