Method serialize_alloc_inputs

exir/serde/serialize.py:166–192

Serialize the inputs to the memory.alloc function. Since there's no specific spec, we jut serialize the inputs with a dummy name. We serialize the AllocSpec into a string "size;dtype"

(
        self, inputs  # pyre-ignore
    )

        return json_acceptable_from_node

    def serialize_alloc_inputs(
        self, inputs  # pyre-ignore
    ) -> List[schema.NamedArgument]:
        """
        Serialize the inputs to the memory.alloc function. Since there's no
        specific spec, we jut serialize the inputs with a dummy name.
        We serialize the AllocSpec into a string "size;dtype"
        """
        assert len(inputs) == 1

        def serialize_alloc_spec(alloc_spec: memory.AllocSpec) -> schema.Argument:
            return schema.Argument.create(
                as_string=f"{alloc_spec[0]};{export_serialize._TORCH_TO_SERIALIZE_DTYPE[alloc_spec[1]].value}"
            )

        if isinstance(inputs[0], list):
            return [
                schema.NamedArgument(name="alloc_list", arg=serialize_alloc_spec(arg))
                for arg in inputs[0]
            ]
        else:
            # Single value
            return [
                schema.NamedArgument(
                    name="alloc_arg", arg=serialize_alloc_spec(inputs[0])
                )
            ]

    def serialize_arbitrary_outputs(self, node: torch.fx.Node) -> List[schema.Argument]:
        meta_val = node.meta["val"]