.torch_interpreter.piece_by_piece_serialize

experimental_experiment.torch_interpreter.piece_by_piece_serialize.choose_kwargs_for_dynamic_shapes(ds_args: Tuple[Dict[int, Any], ...], ds_kwargs: Dict[str, Dict[int, Any]], forward_names: List[str]) → Dict[str, Dict[int, Any]]

Chooses a dictionary to express dynamic shapes when the module uses both unnamed arguments and named ones.

experimental_experiment.torch_interpreter.piece_by_piece_serialize.deserialize_args(res: List[Tensor], expected_types: List[str], clone: bool = False, return_n_args: bool = False) → Tuple[Any, ...]

Deserizalizes output results coming from the custom op and restores the python classes attached to it.

Parameters:

• res – args to deserialize

• expected_types – information on how to deserialize

• clone – clone tensors before returning them

• return_n_args – if True, the function returns the number of deserialized arguments, if False, it assumes this number if equal to the number of expected types

Returns:

new args

experimental_experiment.torch_interpreter.piece_by_piece_serialize.deserialize_args_kwargs(args: List[Tensor], kwargs: Dict[str, Any], expected_types: Tuple[List[str], List[str]], clone: bool = False, ordered_names: List[str] | None = None, fill_kwargs: bool = False) → Tuple[Tuple[Any, ...], Dict[str, Any]]

Deserializes a list of tensor or list of tensors into args and kwargs. kwargs should be empty since this type is allowed as a serialized type.

Parameters:

• args – arguments

• kwargs – named arguments, they should be empty

• expected_types – needed to understand how to deserialize

• clone – clone every tensor

• ordered_names – ordered needed to restore **kwargs

• fill_kwargs – if True, the last parameter is **kwargs and it should be empty

Returns:

new args, new named args

experimental_experiment.torch_interpreter.piece_by_piece_serialize.extract_names_from_schema(schema: str) → List[str]

Extracts name from a C++ schema produced by infer_schema. Example: (Tensor x, Tensor y) -> Tensor returns [“x”, “y”].

experimental_experiment.torch_interpreter.piece_by_piece_serialize.make_copy(obj: Any) → Any

Makes a copy of the objects.

experimental_experiment.torch_interpreter.piece_by_piece_serialize.serialize_args(args: Tuple[Any], kwargs: Dict[str, Any] | None, schema: str, args_names: List[str] | None = None) → Tuple[Tuple[Tensor, ...], Dict[str, Tensor]]

Serializes args and kwargs before calling a custom ops.

Parameters:

• args – unnamed arguments

• kwargs – named arguments

• schema – schema function

• args_names – ordered argument names, it must be specified if kwargs is specified

experimental_experiment.torch_interpreter.piece_by_piece_serialize.serialize_one(obj: Any, name: str | int, schema: str) → Tensor | List[Tensor]

Serializes one object into a tensor or a list of tensors. name and schema are just better error messages.

experimental_experiment.torch_interpreter.piece_by_piece_serialize.tree_spec_as_name(tree_spec: TreeSpec, n_elements: int) → str

Returns a string containing all the information needed to code an instance of TreeSpec in a string.

Parameters:

• tree_spec – instance of TreeSpec to convert into a name

• n_elements – number of elements it serializes

Returns:

string

experimental_experiment.torch_interpreter.piece_by_piece_serialize.tree_spec_from_name(name: str) → Tuple[int, TreeSpec]

Restores the instance of TreeSpec converted into a name by function tree_spec_as_name().

Parameters:

name – name of the TreeSpec

Returns:

instance of TreeSpec, number of elements it contains

experimental_experiment.torch_interpreter.piece_by_piece_serialize.type_as_str_with_info(obj: Any) → str

Returns a string with information about how to deserialize.