def ann_to_type(ann):
if ann is None:
return TensorType.get()
elif ann is torch.Tensor:
return TensorType.get()
elif is_tuple(ann):
return TupleType([ann_to_type(a) for a in ann.__args__])
elif is_list(ann):
return ListType(ann_to_type(ann.__args__[0]))
elif is_dict(ann):
key = ann_to_type(ann.__args__[0])
value = ann_to_type(ann.__args__[1])
return DictType(key, value)
elif is_optional(ann):
if issubclass(ann.__args__[1], type(None)):
return OptionalType(ann_to_type(ann.__args__[0]))
else:
return OptionalType(ann_to_type(ann.__args__[1]))
elif ann is float:
return FloatType.get()
elif ann is int:
return IntType.get()
elif ann is str:
return StringType.get()
elif ann is bool:
return BoolType.get()
elif hasattr(ann, "__torch_script_class__"):
return ClassType(_qualified_name(ann))
raise ValueError("Unknown type annotation: '{}'".format(ann))
def try_ann_to_type(ann, loc):
if ann is None:
return TensorType.get()
if inspect.isclass(ann) and issubclass(ann, torch.Tensor):
return TensorType.get()
if is_tuple(ann):
return TupleType([try_ann_to_type(a, loc) for a in ann.__args__])
if is_list(ann):
elem_type = try_ann_to_type(ann.__args__[0], loc)
if elem_type:
return ListType(elem_type)
if is_dict(ann):
key = try_ann_to_type(ann.__args__[0], loc)
value = try_ann_to_type(ann.__args__[1], loc)
return DictType(key, value)
if is_optional(ann):
if issubclass(ann.__args__[1], type(None)):
return OptionalType(try_ann_to_type(ann.__args__[0], loc))
else:
return OptionalType(try_ann_to_type(ann.__args__[1], loc))
if is_rref(ann):
return RRefType(try_ann_to_type(ann.__args__[0], loc))
if is_future(ann):
return FutureType(try_ann_to_type(ann.__args__[0], loc))
if ann is float:
return FloatType.get()
if ann is int:
return IntType.get()
if ann is str:
return StringType.get()
if ann is bool:
return BoolType.get()
if ann is Any:
return AnyType.get()
if ann is type(None):
return NoneType.get()
if inspect.isclass(ann) and hasattr(ann, "__torch_script_interface__"):
return InterfaceType(_qualified_name(ann))
if ann is torch.device:
return DeviceObjType.get()
if ann is torch.dtype:
return IntType.get() # dtype not yet bound in as its own type
if inspect.isclass(ann):
if hasattr(ann, "__torch_script_class__"):
return ClassType(_qualified_name(ann))
# Why Callable? forward is declared to be a Callable so that
# people can define it without mypy complaining. But we shouldn't
# try to recursively compile it!
ignored_builtin_classes = (torch.nn.Module, tuple, list, Callable)
if torch._jit_internal.can_compile_class(ann) and not issubclass(
ann, ignored_builtin_classes):
torch.jit._recursive_compile_class(ann, loc)
return ClassType(_qualified_name(ann))
# Maybe resolve a NamedTuple to a Tuple Type
def fake_rcb(key):
return None
return torch._C._resolve_type_from_object(ann, loc, fake_rcb)
def try_ann_to_type(ann, loc):
if ann is None:
return TensorType.get()
if inspect.isclass(ann) and issubclass(ann, torch.Tensor):
return TensorType.get()
if is_tuple(ann):
return TupleType([try_ann_to_type(a, loc) for a in ann.__args__])
if is_list(ann):
elem_type = try_ann_to_type(ann.__args__[0], loc)
if elem_type:
return ListType(elem_type)
if is_dict(ann):
key = try_ann_to_type(ann.__args__[0], loc)
value = try_ann_to_type(ann.__args__[1], loc)
return DictType(key, value)
if is_optional(ann):
if issubclass(ann.__args__[1], type(None)):
valid_type = try_ann_to_type(ann.__args__[0], loc)
else:
valid_type = try_ann_to_type(ann.__args__[1], loc)
assert valid_type, "Unsupported annotation {} could not be resolved.".format(
repr(ann))
return OptionalType(valid_type)
if torch.distributed.rpc.is_available() and is_rref(ann):
return RRefType(try_ann_to_type(ann.__args__[0], loc))
if is_future(ann):
return FutureType(try_ann_to_type(ann.__args__[0], loc))
if ann is float:
return FloatType.get()
if ann is int:
return IntType.get()
if ann is str:
return StringType.get()
if ann is bool:
return BoolType.get()
if ann is Any:
return AnyType.get()
if ann is type(None):
return NoneType.get()
if inspect.isclass(ann) and hasattr(ann, "__torch_script_interface__"):
return InterfaceType(_qualified_name(ann))
if ann is torch.device:
return DeviceObjType.get()
if ann is torch.dtype:
return IntType.get() # dtype not yet bound in as its own type
if inspect.isclass(ann):
if hasattr(ann, "__torch_script_class__"):
return ClassType(_qualified_name(ann))
ignored_builtin_classes = (torch.nn.Module, tuple, list)
if torch._jit_internal.can_compile_class(ann) and not issubclass(
ann, ignored_builtin_classes):
torch.jit._script._recursive_compile_class(ann, loc)
return ClassType(_qualified_name(ann))
# Maybe resolve a NamedTuple to a Tuple Type
def fake_rcb(key):
return None
return torch._C._resolve_type_from_object(ann, loc, fake_rcb)
def ann_to_type(ann):
if ann is None:
return DynamicType()
elif ann is torch.Tensor:
return DynamicType()
elif is_tuple(ann):
return TupleType([ann_to_type(a) for a in ann.__args__])
raise ValueError(
"The only supported annotations kinds are Tensor and Tuple[...]")
def try_ann_to_type(ann, loc):
if ann is None:
return TensorType.get()
elif ann is torch.Tensor:
return TensorType.get()
elif is_tuple(ann):
return TupleType([try_ann_to_type(a, loc) for a in ann.__args__])
elif is_list(ann):
elem_type = try_ann_to_type(ann.__args__[0], loc)
if elem_type:
return ListType(elem_type)
elif is_dict(ann):
key = try_ann_to_type(ann.__args__[0], loc)
value = try_ann_to_type(ann.__args__[1], loc)
return DictType(key, value)
elif is_optional(ann):
if issubclass(ann.__args__[1], type(None)):
return OptionalType(try_ann_to_type(ann.__args__[0], loc))
else:
return OptionalType(try_ann_to_type(ann.__args__[1], loc))
elif is_rref(ann):
return RRefType(try_ann_to_type(ann.__args__[0], loc))
elif is_future(ann):
return FutureType(try_ann_to_type(ann.__args__[0], loc))
elif ann is float:
return FloatType.get()
elif ann is int:
return IntType.get()
elif ann is str:
return StringType.get()
elif ann is bool:
return BoolType.get()
elif ann is Any:
return AnyType.get()
elif ann is type(None):
return NoneType.get()
elif inspect.isclass(ann) and hasattr(ann, "__torch_script_class__"):
return ClassType(_qualified_name(ann))
elif inspect.isclass(ann) and hasattr(ann, "__torch_script_interface__"):
return InterfaceType(_qualified_name(ann))
elif ann is torch.device:
return DeviceObjType.get()
else:
# Maybe resolve a NamedTuple to a Tuple Type
def fake_rcb(key):
return None
the_type = torch._C._resolve_type_from_object(ann, loc, fake_rcb)
if the_type is not None:
return the_type
return None
def ann_to_type(ann):
if ann is None:
return DynamicType.get()
elif ann is torch.Tensor:
return DynamicType.get()
elif is_tuple(ann):
return TupleType([ann_to_type(a) for a in ann.__args__])
elif is_list(ann):
return ListType(ann_to_type(ann.__args__[0]))
elif ann is float:
return FloatType.get()
elif ann is int:
return IntType.get()
raise ValueError("The only supported annotations kinds are Tensor and Tuple[...]")
def ann_to_type(ann, resolver=None):
# resolver should be a Tuple[Callable, SourceRange] where the Callable
# is a resolutionCallback
if ann is None:
return TensorType.get()
elif ann is torch.Tensor:
return TensorType.get()
elif is_tuple(ann):
return TupleType([ann_to_type(a) for a in ann.__args__])
elif is_list(ann):
return ListType(ann_to_type(ann.__args__[0]))
elif is_dict(ann):
key = ann_to_type(ann.__args__[0])
value = ann_to_type(ann.__args__[1])
return DictType(key, value)
elif is_optional(ann):
if issubclass(ann.__args__[1], type(None)):
return OptionalType(ann_to_type(ann.__args__[0]))
else:
return OptionalType(ann_to_type(ann.__args__[1]))
elif is_rref(ann):
return RRefType(ann_to_type(ann.__args__[0]))
elif ann is float:
return FloatType.get()
elif ann is int:
return IntType.get()
elif ann is str:
return StringType.get()
elif ann is bool:
return BoolType.get()
elif ann is Any:
return AnyType.get()
elif ann is type(None):
return NoneType.get()
elif hasattr(ann, "__torch_script_class__"):
return ClassType(_qualified_name(ann))
elif hasattr(ann, "__torch_script_interface__"):
return InterfaceType(_qualified_name(ann))
elif ann is torch.device:
return DeviceObjType.get()
elif resolver is not None:
# Maybe resolve a NamedTuple to a Tuple Type
rcb, loc = resolver
the_type = torch._C._resolve_type(ann.__name__, loc, rcb)
if the_type is not None:
return the_type
raise ValueError("Unknown type annotation: '{}'".format(ann))
def ann_to_type(ann):
if ann is None:
return DynamicType.get()
elif ann is torch.Tensor:
return DynamicType.get()
elif is_tuple(ann):
return TupleType([ann_to_type(a) for a in ann.__args__])
elif is_list(ann):
return ListType(ann_to_type(ann.__args__[0]))
elif is_dict(ann):
key = ann_to_type(ann.__args__[0])
value = ann_to_type(ann.__args__[1])
return DictType(key, value)
elif ann is float:
return FloatType.get()
elif ann is int:
return IntType.get()
elif ann is str:
return StringType.get()
raise ValueError("Unknown type annotation: '{}'".format(ann.__name__))
def default_signature(fn, source, _n_arguments, _n_binders):
"""Returns the default signature for fn.
The current formula is to use the source (if available) to determine the
number of inputs and outputs, and set all their types as tensors.
If the source is missing, we fall back to the numbers provided by the compiler,
to make sure we don't cause an error there (although type mismatches can still happen).
This method also accounts for the self argument if fn is a method.
"""
if _n_binders is None:
raise RuntimeError(
"default_signature needs to know the number of binders")
if source is None and _n_arguments is None:
raise RuntimeError(
"default_signature needs either the source or the number of arguments"
)
ret_type = TupleType([DynamicType() for _ in range(_n_binders)])
if source is not None:
py_ast = ast.parse(source)
if len(py_ast.body) != 1 or not isinstance(py_ast.body[0],
ast.FunctionDef):
raise RuntimeError("expected a single top-level function")
py_def = py_ast.body[0]
# TODO: ideally we'd ignore the type of varargs entirely, but we currently don't
# allow passing in anything else than tensors anyway.
if py_def.args.vararg is not None:
arg_types = [DynamicType()] * _n_arguments
else:
arg_types = [DynamicType() for _ in py_def.args.args]
if inspect.ismethod(fn):
arg_types = arg_types[1:]
else:
arg_types = [DynamicType()] * _n_arguments
return arg_types, ret_type
def augment_many_model_functions_with_bundled_inputs(
model: torch.jit.ScriptModule,
inputs: Dict[Callable, Optional[Sequence[Tuple[Any, ...]]]],
_receive_inflate_expr: Optional[List[str]] = None, # For debugging.
info: Optional[Dict[Callable, List[
str]]] = None, # Optional argument to provide info about the function or its inputs
skip_size_check=False,
) -> None:
"""Add bundled sample inputs to a model for an arbitrary list of public functions.
Models with bundled inputs can be invoked in a uniform manner by
benchmarking and code coverage tools.
Augmented models will support the following methods:
`get_all_bundled_inputs_for_<function_name>() -> List[Tuple[Any, ...]]`
Returns a list of tuples suitable for passing to the model like
`for inp in model.get_all_bundled_inputs_for_foo(): model.foo(*inp)`
`get_bundled_inputs_functions_and_info() -> Dict[str, Dict[str: List[str]]]`
Returns a dictionary mapping function names to a metadata dictionary.
This nested dictionary maps preset strings like:
'get_inputs_function_name' -> the name of a function attribute in this model that can be
run to get back a list of inputs corresponding to that function.
'info' -> the user provided extra information about the bundled inputs
If forward has bundled inputs then these following functions are also defined:
`get_all_bundled_inputs() -> List[Tuple[Any, ...]]`
Returns a list of tuples suitable for passing to the model like
`for inp in model.get_all_bundled_inputs(): model(*inp)`
`get_num_bundled_inputs() -> int`
Equivalent to `len(model.get_all_bundled_inputs())`,
but slightly easier to call from C++.
Inputs can be specified in one of two ways:
- The model can define `_generate_bundled_inputs_for_<function_name>`.
If the user chooses this method inputs[<function>] should map to None
- The `inputs` argument to this function can be a dictionary mapping functions to a
list of inputs, of the same form that will be returned by get_all_bundled_inputs_for_<function_name>.
The type of the inputs is List[Tuple[Any, ...]]. The outer list corresponds with a
list of inputs, the inner tuple is the list of args that together make up one input.
For inputs of functions that take one arg, this will be a tuple of length one. The Any, ...
is the actual data that makes up the args, e.g. a tensor.
Info is an optional parameter that maps functions to a list of strings providing extra information about that
function's bundled inputs. This could be descriptions, expected outputs, etc.
- Ex: info={model.forward : ['man eating icecream', 'an airplane', 'a dog']}
This function will attempt to optimize arguments so that (e.g.)
arguments like `torch.zeros(1000)` will be represented compactly.
Only top-level arguments will be optimized.
Tensors in lists or tuples will not.
"""
if not isinstance(model, torch.jit.ScriptModule):
raise Exception("Only ScriptModule is supported.")
if not inputs:
raise Exception("Please provide inputs for at least 1 function")
if hasattr(model, "get_all_bundled_inputs") or hasattr(
model, "get_bundled_inputs_functions_and_info"):
raise Exception(
"Models can only be augmented with bundled inputs once. "
"This Model seems to have already been augmented with "
"bundled inputs. Please start afresh with one that "
"doesn't have bundled inputs.", )
get_bundled_inputs_functions_and_info_template = ""
for function, input_list in inputs.items():
if hasattr(function, "__name__"):
function_name = function.__name__
else:
if hasattr(function, "name"):
function_name = function.name # type: ignore[attr-defined]
else:
raise Exception(
'At least one of your functions has no attribute name please ensure all have one. m.foo.name = "foo"'
)
if input_list is not None and not isinstance(input_list, Sequence):
raise TypeError(
"Error inputs for function {0} is not a Sequence".format(
function_name))
function_arg_types = [
arg.type for arg in function.schema.arguments[1:]
] # type: ignore[attr-defined]
deflated_inputs_type: ListType = ListType(
TupleType(function_arg_types))
model._c._register_attribute(
"_bundled_inputs_deflated_{name}".format(name=function_name),
deflated_inputs_type, [])
if hasattr(model, "_generate_bundled_inputs_for_" + function_name):
if input_list is not None:
raise Exception(
"inputs[{name}] is not None, but _generate_bundled_inputs_for_{name} is already defined"
.format(name=function_name))
# Model author already defined _generate_bundled_inputs_for_<function_name>.
elif input_list is None or len(input_list) == 0:
raise Exception(
"inputs for {name} must be specified if _generate_bundled_inputs_for_{name} is not already defined"
.format(name=function_name, ))
else:
# Iterate over the inputs and args in each input.
# Accumulate `deflated_inputs` as (possibly) compressed values
# and `parts` to be joined into the expression that unpacks them.
deflated_inputs = []
parts = []
for inp_idx, args in enumerate(input_list):
if not isinstance(args, Tuple) and not isinstance(
args, List): # type: ignore[arg-type]
raise TypeError(
"Error bundled input for function {0} idx: {1} is not a Tuple or a List"
.format(function_name, inp_idx))
deflated_args = []
parts.append("(")
for arg_idx, arg in enumerate(args):
inflate_helper_fn_name = _get_inflate_helper_fn_name(
arg_idx, inp_idx, function_name)
deflated, inflater, helper_definition = _inflate_expr(
arg,
f"deflated[{inp_idx}][{arg_idx}]",
inflate_helper_fn_name,
skip_size_check=skip_size_check,
)
deflated_args.append(deflated)
parts.append(f" {inflater},")
if helper_definition:
model.define(textwrap.dedent(helper_definition))
deflated_inputs.append(tuple(deflated_args))
parts.append("),")
parts.append("")
expr = "\n".join(parts)
# Back-channel return this expr for debugging.
if _receive_inflate_expr is not None:
_receive_inflate_expr.append(expr)
setattr(
model,
"_bundled_inputs_deflated_{name}".format(name=function_name),
deflated_inputs)
definition = textwrap.dedent("""
def _generate_bundled_inputs_for_{name}(self):
deflated = self._bundled_inputs_deflated_{name}
return [
{expr}
]
""").format(expr=expr, name=function_name)
model.define(definition)
# Define get_all_bundled_inputs_for_<function_name> that caches the generated inputs.
model.define(
textwrap.dedent("""
def get_all_bundled_inputs_for_{name}(self):
all_inputs = self._generate_bundled_inputs_for_{name}()
assert all_inputs is not None
return all_inputs
""").format(name=function_name))
# Add to the high level helper methods
inputs_info = repr(
info[function]) if info and function in info else '[]'
get_bundled_inputs_functions_and_info_template += """
temp_dict : Dict[str,List[str]] = {{}}
info: List[str] = {info}
temp_dict['info'] = info
temp_dict['get_inputs_function_name'] = ['get_all_bundled_inputs_for_{name}']
all_inputs['{name}'] = temp_dict
""".format(
name=function_name,
info=inputs_info,
)
# To ensure backwards compatibility and a streamlined api for forward these wrappers are provided
if function_name == 'forward':
model.define(
textwrap.dedent("""
def get_all_bundled_inputs(self):
return self.get_all_bundled_inputs_for_forward()
"""))
model.define(
textwrap.dedent("""
def get_num_bundled_inputs(self):
return len(self.get_all_bundled_inputs_for_forward())
"""))
# Define some high level helper methods that act on all bundled inputs
model.define(
textwrap.dedent("""
def get_bundled_inputs_functions_and_info(self):
all_inputs : Dict[str, Dict[str,List[str]]] = {{}}
{template}
return all_inputs
""".format(template=get_bundled_inputs_functions_and_info_template)))
def try_ann_to_type(ann, loc):
if ann is inspect.Signature.empty:
return TensorType.getInferred()
if ann is None:
return NoneType.get()
if inspect.isclass(ann) and is_tensor(ann):
return TensorType.get()
if is_tuple(ann):
# Special case for the empty Tuple type annotation `Tuple[()]`
if len(ann.__args__) == 1 and ann.__args__[0] == ():
return TupleType([])
return TupleType([try_ann_to_type(a, loc) for a in ann.__args__])
if is_list(ann):
elem_type = try_ann_to_type(ann.__args__[0], loc)
if elem_type:
return ListType(elem_type)
if is_dict(ann):
key = try_ann_to_type(ann.__args__[0], loc)
value = try_ann_to_type(ann.__args__[1], loc)
# Raise error if key or value is None
if key is None:
raise ValueError(f"Unknown type annotation: '{ann.__args__[0]}' at {loc.highlight()}")
if value is None:
raise ValueError(f"Unknown type annotation: '{ann.__args__[1]}' at {loc.highlight()}")
return DictType(key, value)
if is_optional(ann):
if issubclass(ann.__args__[1], type(None)):
contained = ann.__args__[0]
else:
contained = ann.__args__[1]
valid_type = try_ann_to_type(contained, loc)
msg = "Unsupported annotation {} could not be resolved because {} could not be resolved."
assert valid_type, msg.format(repr(ann), repr(contained))
return OptionalType(valid_type)
if is_union(ann):
# TODO: this is hack to recognize NumberType
if set(ann.__args__) == set([int, float, complex]):
return NumberType.get()
inner: List = []
# We need these extra checks because both `None` and invalid
# values will return `None`
# TODO: Determine if the other cases need to be fixed as well
for a in ann.__args__:
if a is None:
inner.append(NoneType.get())
maybe_type = try_ann_to_type(a, loc)
msg = "Unsupported annotation {} could not be resolved because {} could not be resolved."
assert maybe_type, msg.format(repr(ann), repr(maybe_type))
inner.append(maybe_type)
return UnionType(inner) # type: ignore[arg-type]
if torch.distributed.rpc.is_available() and is_rref(ann):
return RRefType(try_ann_to_type(ann.__args__[0], loc))
if is_future(ann):
return FutureType(try_ann_to_type(ann.__args__[0], loc))
if ann is float:
return FloatType.get()
if ann is complex:
return ComplexType.get()
if ann is int:
return IntType.get()
if ann is str:
return StringType.get()
if ann is bool:
return BoolType.get()
if ann is Any:
return AnyType.get()
if ann is type(None):
return NoneType.get()
if inspect.isclass(ann) and hasattr(ann, "__torch_script_interface__"):
return InterfaceType(ann.__torch_script_interface__)
if ann is torch.device:
return DeviceObjType.get()
if ann is torch.Stream:
return StreamObjType.get()
if ann is torch.dtype:
return IntType.get() # dtype not yet bound in as its own type
if inspect.isclass(ann) and issubclass(ann, enum.Enum):
if _get_script_class(ann) is None:
scripted_class = torch.jit._script._recursive_compile_class(ann, loc)
name = scripted_class.qualified_name()
else:
name = _qualified_name(ann)
return EnumType(name, get_enum_value_type(ann, loc), list(ann))
if inspect.isclass(ann):
maybe_script_class = _get_script_class(ann)
if maybe_script_class is not None:
return maybe_script_class
if torch._jit_internal.can_compile_class(ann):
return torch.jit._script._recursive_compile_class(ann, loc)
# Maybe resolve a NamedTuple to a Tuple Type
def fake_rcb(key):
return None
return torch._C._resolve_type_from_object(ann, loc, fake_rcb)
def try_ann_to_type(ann, loc):
if ann is None:
return TensorType.getInferred()
if inspect.isclass(ann) and issubclass(ann, torch.Tensor):
return TensorType.get()
if is_tuple(ann):
return TupleType([try_ann_to_type(a, loc) for a in ann.__args__])
if is_list(ann):
elem_type = try_ann_to_type(ann.__args__[0], loc)
if elem_type:
return ListType(elem_type)
if is_dict(ann):
key = try_ann_to_type(ann.__args__[0], loc)
value = try_ann_to_type(ann.__args__[1], loc)
# Raise error if key or value is None
if key is None:
raise ValueError(
f"Unknown type annotation: '{ann.__args__[0]}' at {loc.highlight()}"
)
if value is None:
raise ValueError(
f"Unknown type annotation: '{ann.__args__[1]}' at {loc.highlight()}"
)
return DictType(key, value)
if is_optional(ann):
if issubclass(ann.__args__[1], type(None)):
contained = ann.__args__[0]
else:
contained = ann.__args__[1]
valid_type = try_ann_to_type(contained, loc)
msg = "Unsupported annotation {} could not be resolved because {} could not be resolved."
assert valid_type, msg.format(repr(ann), repr(contained))
return OptionalType(valid_type)
if torch.distributed.rpc.is_available() and is_rref(ann):
return RRefType(try_ann_to_type(ann.__args__[0], loc))
if is_future(ann):
return FutureType(try_ann_to_type(ann.__args__[0], loc))
if ann is float:
return FloatType.get()
if ann is complex:
return ComplexType.get()
if ann is int:
return IntType.get()
if ann is str:
return StringType.get()
if ann is bool:
return BoolType.get()
if ann is Any:
return AnyType.get()
if ann is type(None):
return NoneType.get()
if inspect.isclass(ann) and hasattr(ann, "__torch_script_interface__"):
return InterfaceType(_qualified_name(ann))
if ann is torch.device:
return DeviceObjType.get()
if ann is torch.Stream:
return StreamObjType.get()
if ann is torch.dtype:
return IntType.get() # dtype not yet bound in as its own type
if inspect.isclass(ann) and issubclass(ann, enum.Enum):
qualified_name = _qualified_name(ann)
if _get_script_class(qualified_name) is None:
torch.jit._script._recursive_compile_class(ann, loc)
return EnumType(_qualified_name(ann), get_enum_value_type(ann, loc),
list(ann))
if inspect.isclass(ann):
qualified_name = _qualified_name(ann)
if _get_script_class(qualified_name) is not None:
return ClassType(qualified_name)
ignored_builtin_classes = (torch.nn.Module, tuple, list, Exception)
if torch._jit_internal.can_compile_class(ann) and not issubclass(
ann, ignored_builtin_classes):
torch.jit._script._recursive_compile_class(ann, loc)
return ClassType(qualified_name)
# Maybe resolve a NamedTuple to a Tuple Type
def fake_rcb(key):
return None
return torch._C._resolve_type_from_object(ann, loc, fake_rcb)
def augment_model_with_bundled_inputs(
model: torch.jit.ScriptModule,
inputs: Optional[List[Tuple[Any, ...]]] = None,
_receive_inflate_expr: Optional[List[str]] = None, # For debugging.
) -> None:
"""Add bundled sample inputs to a model.
Models with bundled inputs can be invoked in a uniform manner by
benchmarking and code coverage tools.
Augmented models will support the following methods:
`get_all_bundled_inputs() -> List[Tuple[Any, ...]]`
Returns a list of tuples suitable for passing to the model like
`for inp in model.get_all_bundled_inputs(): model(*inp)`
`get_num_bundled_inputs() -> int`
Equivalent to `len(model.get_all_bundled_inputs())`,
but slightly easier to call from C++.
`run_on_bundled_input(idx: int) -> Any`
Run the model on bundled input number `idx`
Inputs can be specified in one of two ways:
- The model can define `_generate_bundled_inputs`
get_all_bundled_inputs will simply call this method
and cache the value.
- The `inputs` argument to this function can be a list of tuples,
of the same form that will be returned by get_all_bundled_inputs.
This function will attempt to optimize arguments so that (e.g.)
arguments like `torch.zeros(1000)` will be represented compactly.
Only top-level arguments will be optimized.
Tensors in lists or tuples will not.
"""
if not isinstance(model, torch.jit.ScriptModule):
raise Exception("Only ScriptModule is supported.")
forward_arg_types = [
arg.type for arg in model.forward.schema.arguments[1:]
]
deflated_inputs_type: ListType = ListType(TupleType(forward_arg_types))
inflated_inputs_type: OptionalType[ListType] = OptionalType(
deflated_inputs_type)
model._c._register_attribute("_bundled_inputs_deflated",
deflated_inputs_type, [])
model._c._register_attribute("_bundled_inputs_inflated",
inflated_inputs_type, None)
if hasattr(model, "_generate_bundled_inputs"):
if inputs is not None:
raise Exception(
"inputs is not None, but _generate_bundled_inputs is already defined"
)
# Model author already defined _generate_bundled_inputs.
elif inputs is None:
raise Exception(
"inputs must be specified if _generate_bundled_inputs is not already defined"
)
else:
# Iterate over the inputs and args in each input.
# Accumulate `deflated_inputs` as (possibly) compressed values
# and `parts` to be joined into the expression that unpacks them.
deflated_inputs = []
parts = []
for inp_idx, args in enumerate(inputs):
deflated_args = []
parts.append("(")
for arg_idx, arg in enumerate(args):
deflated, inflater = _inflate_expr(
arg, f"deflated[{inp_idx}][{arg_idx}]")
deflated_args.append(deflated)
parts.append(f" {inflater},")
deflated_inputs.append(tuple(deflated_args))
parts.append("),")
parts.append("")
expr = "\n".join(parts)
# Back-channel return this expr for debugging.
if _receive_inflate_expr is not None:
_receive_inflate_expr.append(expr)
model._bundled_inputs_deflated = deflated_inputs
definition = textwrap.dedent("""
def _generate_bundled_inputs(self):
deflated = self._bundled_inputs_deflated
return [
{}
]
""").format(expr)
model.define(definition)
# Define get_all_bundled_inputs that caches the generated inputs.
model.define(
textwrap.dedent("""
def get_all_bundled_inputs(self):
if self._bundled_inputs_inflated is None:
self._bundled_inputs_inflated = self._generate_bundled_inputs()
all_inputs = self._bundled_inputs_inflated
assert all_inputs is not None
return all_inputs
"""))
# Define some helper methods.
model.define(
textwrap.dedent("""
def get_num_bundled_inputs(self):
return len(self.get_all_bundled_inputs())
"""))
model.define(
textwrap.dedent("""
def run_on_bundled_input(self, idx: int):
return self(*self.get_all_bundled_inputs()[idx])
"""))
def augment_many_model_functions_with_bundled_inputs(
model: torch.jit.ScriptModule,
inputs: Dict[Callable, Optional[Sequence[Tuple[Any, ...]]]],
_receive_inflate_expr: Optional[List[str]] = None, # For debugging.
info: Optional[Dict[Callable, List[
str]]] = None, # Optional argument to provide info about the function or its inputs
) -> None:
"""Add bundled sample inputs to a model for an arbitrary list of public functions.
Models with bundled inputs can be invoked in a uniform manner by
benchmarking and code coverage tools.
Augmented models will support the following methods:
`get_all_bundled_inputs_for_<function_name>() -> List[Tuple[Any, ...]]`
Returns a list of tuples suitable for passing to the model like
`for inp in model.get_all_bundled_inputs_for_foo(): model.foo(*inp)`
`get_bundled_inputs_functions_and_info() -> Dict[str, Dict[str: List[str]]]`
Returns a dictionary mapping function names to a metadata dictionary.
This nested dictionary maps preset strings like:
'get_inputs_function_name' -> the name of a function attribute in this model that can be
run to get back a list of inputs corresponding to that function.
'info' -> the user provided extra information about the bundled inputs
If forward has bundled inputs then these following functions are also defined:
`get_all_bundled_inputs() -> List[Tuple[Any, ...]]`
Returns a list of tuples suitable for passing to the model like
`for inp in model.get_all_bundled_inputs(): model(*inp)`
`get_num_bundled_inputs() -> int`
Equivalent to `len(model.get_all_bundled_inputs())`,
but slightly easier to call from C++.
`run_on_bundled_input(idx: int) -> Any`
Run the model on bundled input number `idx`
Inputs can be specified in one of two ways:
- The model can define `_generate_bundled_inputs_for_<function_name>`
get_all_bundled_inputs will simply call this method
and cache the value. If the user chooses this method inputs[<function>]
should map to None
- The `inputs` argument to this function can be a dictionary mapping functions to a
list of tuples, of the same form that will be returned by get_all_bundled_inputs_for_<function_name>.
It is highly recommended (though not enforced) that if multiple functions have the same input style, that
you create separate bundled inputs for each function. Reusing the same input and bundling it to multiple
functions can cause issues with other torch.jit functionality like freeze
Info is an optional parameter that maps functions to a list of strings providing extra information about that
function's bundled inputs. This could be descriptions, expected outputs, etc.
- Ex: info={model.forward : ['man eating icecream', 'an airplane', 'a dog']}
This function will attempt to optimize arguments so that (e.g.)
arguments like `torch.zeros(1000)` will be represented compactly.
Only top-level arguments will be optimized.
Tensors in lists or tuples will not.
"""
if not isinstance(model, torch.jit.ScriptModule):
raise Exception("Only ScriptModule is supported.")
get_bundled_inputs_functions_and_info_template = ""
for function, input_list in inputs.items():
function_name = function.__name__
function_arg_types = [
arg.type for arg in function.schema.arguments[1:]
] # type: ignore
deflated_inputs_type: ListType = ListType(
TupleType(function_arg_types))
inflated_inputs_type: OptionalType[ListType] = OptionalType(
deflated_inputs_type)
model._c._register_attribute(
"_bundled_inputs_deflated_{name}".format(name=function_name),
deflated_inputs_type, [])
model._c._register_attribute(
"_bundled_inputs_inflated_{name}".format(name=function_name),
inflated_inputs_type, None)
if hasattr(model, "_generate_bundled_inputs_for_" + function_name):
if input_list is not None:
raise Exception(
"inputs[{name}] is not None, but _generate_bundled_inputs_for_{name} is already defined"
.format(name=function_name))
# Model author already defined _generate_bundled_inputs_for_<function_name>.
elif input_list is None or len(input_list) == 0:
raise Exception(
"inputs for {name} must be specified if _generate_bundled_inputs_for_{name} is not already defined"
.format(name=function_name, ))
else:
# Iterate over the inputs and args in each input.
# Accumulate `deflated_inputs` as (possibly) compressed values
# and `parts` to be joined into the expression that unpacks them.
deflated_inputs = []
parts = []
for inp_idx, args in enumerate(input_list):
deflated_args = []
parts.append("(")
for arg_idx, arg in enumerate(args):
deflated, inflater = _inflate_expr(
arg, f"deflated[{inp_idx}][{arg_idx}]")
deflated_args.append(deflated)
parts.append(f" {inflater},")
deflated_inputs.append(tuple(deflated_args))
parts.append("),")
parts.append("")
expr = "\n".join(parts)
# Back-channel return this expr for debugging.
if _receive_inflate_expr is not None:
_receive_inflate_expr.append(expr)
model._bundled_inputs_deflated = deflated_inputs
setattr(
model,
"_bundled_inputs_deflated_{name}".format(name=function_name),
deflated_inputs)
definition = textwrap.dedent("""
def _generate_bundled_inputs_for_{name}(self):
deflated = self._bundled_inputs_deflated_{name}
return [
{expr}
]
""").format(expr=expr, name=function_name)
model.define(definition)
# Define get_all_bundled_inputs_for_<function_name> that caches the generated inputs.
model.define(
textwrap.dedent("""
def get_all_bundled_inputs_for_{name}(self):
if self._bundled_inputs_inflated_{name} is None:
self._bundled_inputs_inflated_{name} = self._generate_bundled_inputs_for_{name}()
all_inputs = self._bundled_inputs_inflated_{name}
assert all_inputs is not None
return all_inputs
""").format(name=function_name))
# Add to the high level helper methods
inputs_info = repr(
info[function]) if info and function in info else '[]'
get_bundled_inputs_functions_and_info_template += """
temp_dict : Dict[str,List[str]] = {{}}
info: List[str] = {info}
temp_dict['info'] = info
temp_dict['get_inputs_function_name'] = ['get_all_bundled_inputs_for_{name}']
all_inputs['{name}'] = temp_dict
""".format(
name=function_name,
info=inputs_info,
)
# To ensure backwards compatibility and a streamlined api for forward these wrappers are provided
if function_name == 'forward':
model.define(
textwrap.dedent("""
def get_all_bundled_inputs(self):
return self.get_all_bundled_inputs_for_forward()
"""))
model.define(
textwrap.dedent("""
def get_num_bundled_inputs(self):
return len(self.get_all_bundled_inputs_for_forward())
"""))
model.define(
textwrap.dedent("""
def run_on_bundled_input(self, idx: int):
return self(*self.get_all_bundled_inputs()[idx])
"""))
# Define some high level helper methods that act on all bundled inputs
model.define(
textwrap.dedent("""
def get_bundled_inputs_functions_and_info(self):
all_inputs : Dict[str, Dict[str,List[str]]] = {{}}
{template}
return all_inputs
""".format(template=get_bundled_inputs_functions_and_info_template)))