tags: set[str],
      serialized_signature_def_map: dict[str, bytes],
  ) -> Optional[bool]: ...
  # LINT.ThenChange()

  # LINT.IfChange(run_calibration)
  def run_calibration(
      self,
      saved_model_path: str,
      signature_keys: list[str],
      tags: set[str],
      force_graph_mode_calibration: bool,
      # Value type: RepresentativeDatasetFile.
      representative_dataset_file_map_serialized: dict[str, bytes],

          signature_def_map) const = 0;
  // LINT.ThenChange(
  //     pywrap_function_lib.pyi:save_exported_model,
  //     py_function_lib.py:save_exported_model,
  // )

  // Runs calibration on a model saved at `saved_model_path`. `exported_model`
  // should be the corresponding exported model resulting from the
  // pre-calibration step. `signature_keys` is a set of keys that identify a

        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

    """
    for _ in range(num_examples):
      yield {input_key: random_ops.random_uniform(shape, minval, maxval, dtype)}

  def _save_tf1_model(
      self,
      sess: session.Session,
      saved_model_path: str,
      signature_key: str,
      tags: Collection[str],
      inputs: Mapping[str, core.Tensor],
      outputs: Mapping[str, core.Tensor],
      init_op: Optional[ops.Operation] = None,

_SignatureDefMap = Mapping[str, meta_graph_pb2.SignatureDef]


def get_signatures_from_saved_model(
    saved_model_path: str,
    signature_keys: Optional[Sequence[str]] = None,
    tags: Optional[Collection[str]] = None,
) -> Dict[str, meta_graph_pb2.SignatureDef]:
  """Gets a map from signature keys to their SignatureDef.

  Args:
    saved_model_path: Path to the saved model.

    # (specified by `kDefaultConstantSizeThresholdInBytes`) for unfreezing.
    self._create_matmul_model(
        input_shape=(1, 20),
        weight_shape=(20, 4096),
        saved_model_path=self._input_saved_model_path,
    )

    repr_ds = self._create_data_generator(
        input_key='input_tensor', shape=(1, 20), num_examples=2
    )

    tags = {tag_constants.SERVING}

      the function matching the signature key.

      Raises `StatusNotOk` exception if when the run was unsuccessful.
      )pbdoc",
      py::arg("saved_model_path"), py::arg("dst_saved_model_path"),
      py::arg("quantization_options_serialized"), py::kw_only(),
      py::arg("signature_keys"), py::arg("signature_def_map_serialized"),
      py::arg("py_function_library"),

        ),
        1 / 3,
    )

  def test_when_preset_not_srq_raises_error(self):
    self._create_matmul_model(
        input_shape=(1, 1024),
        weight_shape=(1024, 3),
        saved_model_path=self._input_saved_model_path,
    )

    config = qc.QuantizationConfig()
    with self.assertRaisesRegex(ValueError, 'only supports static-range PTQ'):
      quantization.quantize_saved_model(

    if FLAGS.save_model_path:
      save_model_path = FLAGS.save_model_path
    else:
      save_model_path = tempfile.mkdtemp(suffix='.saved_model')
    save_options = tf.saved_model.SaveOptions(save_debug_info=show_debug_info)
    tf.saved_model.save(
        create_module_fn(), save_model_path, options=save_options
    )
    logging.info('Saved model to: %s', save_model_path)