MCPcopy Create free account
hub / github.com/NVIDIA/TensorRT / export_impl

Function export_impl

tools/Polygraphy/polygraphy/mod/exporter.py:110–207  ·  view source on GitHub ↗
(func_or_cls)

Source from the content-addressed store, hash-verified

108 assert False, f"Could not find method: {method} in the inheritance hierarcy of: {symbol}"
109
110 def export_impl(func_or_cls):
111 _add_to_all(func_or_cls.__name__, module)
112
113 if funcify:
114 # We only support funcify-ing BaseLoaders, and only if __init__ and call_impl
115 # have no overlapping parameters.
116 from polygraphy.backend.base import BaseLoader
117
118 assert inspect.isclass(func_or_cls), "Decorated type must be a loader to use funcify=True"
119 assert BaseLoader in inspect.getmro(
120 func_or_cls
121 ), "Decorated type must derive from BaseLoader to use funcify=True"
122
123 def get_params(method):
124 return list(inspect.signature(find_method(func_or_cls, method)).parameters.values())[1:]
125
126 def is_variadic(param):
127 return param.kind in [param.VAR_POSITIONAL, param.VAR_KEYWORD]
128
129 def has_default(param):
130 return param.default != param.empty
131
132 def get_param_name(p):
133 # For variadic arguments, p.name will drop the *, **
134 return str(p) if is_variadic(p) else p.name
135
136 def param_names(params):
137 return [get_param_name(p) for p in params]
138
139 loader = func_or_cls
140
141 init_params = get_params("__init__")
142 call_impl_params = get_params("call_impl")
143
144 assert (set(param_names(call_impl_params)) - set(param_names(init_params))) == set(
145 param_names(call_impl_params)
146 ), "Cannot funcify a type where call_impl and __init__ have the same argument names!"
147
148 # Dynamically generate a function with the right signature.
149
150 # To generate the signature, we use the init and call_impl arguments,
151 # but move required arguments (i.e. without default values) to the front.
152
153 def build_arg_list(should_include):
154 def str_from_param(p):
155 return get_param_name(p) + (f"={p.default}" if has_default(p) else "")
156
157 arg_list = [str_from_param(p) for p in init_params if should_include(p)]
158 arg_list += [str_from_param(p) for p in call_impl_params if should_include(p)]
159 return arg_list
160
161 non_default_args = build_arg_list(should_include=lambda p: not is_variadic(p) and not has_default(p))
162 default_args = build_arg_list(should_include=lambda p: not is_variadic(p) and has_default(p))
163 special_args = build_arg_list(should_include=is_variadic)
164
165 signature = ", ".join(non_default_args + default_args + special_args)
166
167 init_args = ", ".join(param_names(init_params))

Callers

nothing calls this directly

Calls 9

_add_to_allFunction · 0.85
get_paramsFunction · 0.85
param_namesFunction · 0.85
build_arg_listFunction · 0.85
is_variadicFunction · 0.85
has_defaultFunction · 0.85
pascal_to_snakeFunction · 0.85
try_add_method_docFunction · 0.85
_define_in_moduleFunction · 0.85

Tested by

no test coverage detected