from tensorflow.core.protobuf import meta_graph_pb2

# Mapping of signature def key -> SignatureDef.
_SignatureDefMap = Mapping[str, meta_graph_pb2.SignatureDef]


def _serialize_signature_def_map(
    signature_def_map: _SignatureDefMap,
) -> dict[str, bytes]:
  """Serializes SignatureDef values in `signature_def_map`.

  Args:
    signature_def_map: Signature key -> SignatureDef mapping.

    ):
      return True
    for function in meta_graph.graph_def.library.function:
      if any(node.op.startswith('FakeQuant') for node in function.node_def):
        return True
  return False


def _serialize_signature_def_map(
    signature_def_map: _SignatureDefMap,
) -> dict[str, bytes]:
  """Serializes SignatureDef values in `signature_def_map`.

  Args:
    signature_def_map: Signature key -> SignatureDef mapping.

  def save_exported_model(
      self,
      dst_saved_model_path: str,
      exported_model_serialized: bytes,
      src_saved_model_path: str,
      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],

           py::arg("dst_saved_model_path"),
           py::arg("exported_model_serialized"),
           py::arg("src_saved_model_path"), py::arg("tags"),
           py::arg("serialized_signature_def_map"))
      .def("run_calibration", &PyFunctionLibrary::RunCalibration,
           py::arg("saved_model_path"), py::arg("signature_keys"),
           py::arg("tags"), py::arg("force_graph_mode_calibration"),

  def save_exported_model(
      self,
      dst_saved_model_path: str,
      exported_model_serialized: bytes,
      src_saved_model_path: str,
      tags: set[str],
      serialized_signature_def_map: dict[str, bytes],
  ) -> Optional[bool]:
    # LINT.ThenChange(py_function_lib.h:save_exported_model)
    """Saves `ExportedModel` to `dst_saved_model_path` as a SavedModel.

    Args: