Method serialize_lowered_module

exir/serde/serialize.py:263–308

(self, lowered_module_arg: torch.fx.Node)

        return serialized_args

    def serialize_lowered_module(self, lowered_module_arg: torch.fx.Node) -> str:
        assert lowered_module_arg.op == "get_attr"
        assert isinstance(lowered_module_arg.target, str)

        def serialize_bytes(b: bytes) -> str:
            # We want to serialize the bytes to string because JSON cannot
            # serialize bytes.
            # Since the given bytes may be serialized with any encoding, so we
            # want to first encode with base64, and then decode it with
            # ascii. During deserialization we can just directly decode with b64
            # to get the original encoded bytes.
            return base64.b64encode(b).decode("ascii")

        lowered_module = getattr(
            lowered_module_arg.graph.owning_module, lowered_module_arg.target
        )
        assert isinstance(lowered_module, ExirLoweredBackendModule)

        serialized_compile_spec = [
            CompileSpec(cs.key, serialize_bytes(cs.value))
            for cs in lowered_module.compile_specs
        ]

        serialized_artifact = ExportedProgramSerializer().serialize(
            lowered_module.original_module
        )
        assert isinstance(serialized_artifact.exported_program, schema.ExportedProgram)

        serialized_processed_bytes = serialize_bytes(lowered_module.processed_bytes)
        named_data_store = json.dumps(export_serialize._dataclass_to_dict(lowered_module.named_data_store_output),cls=export_serialize.EnumEncoder) if lowered_module.named_data_store_output else None

        serialized_lowered_module = SerdeLoweredBackendModule(
            original_module=serialized_artifact.exported_program,
            original_state_dict=serialize_bytes(serialized_artifact.state_dict),
            original_constants=serialize_bytes(serialized_artifact.constants),
            processed_bytes=serialized_processed_bytes,
            compile_specs=serialized_compile_spec,
            backend_id=lowered_module.backend_id,
            named_data_store=named_data_store,
        )

        json_lowered_module = json.dumps(
            export_serialize._dataclass_to_dict(serialized_lowered_module),
            cls=export_serialize.EnumEncoder,
        )
        return json_lowered_module


class ExportedProgramSerializer(export_serialize.ExportedProgramSerializer):