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 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):
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