def _get_bundled_inputs_preserved_attributes(
script_module: torch.jit.ScriptModule,
preserved_methods: List[str]) -> List[str]:
bundled_inputs_attributes = []
# Has bundled inputs for forward
if hasattr(script_module, 'get_all_bundled_inputs'):
bundled_inputs_attributes.append('get_all_bundled_inputs')
bundled_inputs_attributes.append('get_num_bundled_inputs')
bundled_inputs_attributes.append('run_on_bundled_input')
# Bundled inputs in module after the change that introduced bundled inputs for multiple functions
if hasattr(script_module, 'get_bundled_inputs_functions_and_info'):
bundled_inputs_attributes.append(
'get_bundled_inputs_functions_and_info')
all_info = script_module.get_bundled_inputs_functions_and_info()
for function_name in all_info:
if function_name not in preserved_methods:
bundled_inputs_attributes.append(function_name)
bundled_inputs_attributes.append("get_all_bundled_inputs_for_" +
function_name)
bundled_inputs_attributes.append("_bundled_inputs_deflated_" +
function_name)
return bundled_inputs_attributes
def _get_bundled_inputs_preserved_attributes(script_module: torch.jit.ScriptModule, preserved_methods: List[str]) -> List[str]:
# Technically it is possible that if a function only bundles inputs for functions besides forward that these wont exist.
# Haven't seen a reason for that to be a valid usecase yet so not going to account for it
bundled_inputs_attributes = [
'get_all_bundled_inputs',
'get_num_bundled_inputs',
'run_on_bundled_input',
]
if hasattr(script_module, 'get_bundled_inputs_functions_and_info'):
bundled_inputs_attributes.append('get_bundled_inputs_functions_and_info')
all_info = script_module.get_bundled_inputs_functions_and_info()
for function_name in all_info:
if function_name not in preserved_methods:
bundled_inputs_attributes.append(function_name)
bundled_inputs_attributes.append("get_all_bundled_inputs_for_" + function_name)
bundled_inputs_attributes.append("_bundled_inputs_deflated_" + function_name)
return bundled_inputs_attributes
def _get_bundled_inputs_attributes_and_methods(
script_module: torch.jit.ScriptModule) -> Tuple[List[str], List[str]]:
methods: List[str] = []
attributes: List[str] = []
# Has bundled inputs for forward
if hasattr(script_module, 'get_all_bundled_inputs'):
methods.append('get_all_bundled_inputs')
methods.append('get_num_bundled_inputs')
methods.append('run_on_bundled_input')
if hasattr(script_module, 'get_bundled_inputs_functions_and_info'):
methods.append('get_bundled_inputs_functions_and_info')
all_info = script_module.get_bundled_inputs_functions_and_info()
for function_name in all_info:
methods.append("get_all_bundled_inputs_for_" + function_name)
methods.append("_generate_bundled_inputs_for_" + function_name)
attributes.append("_bundled_inputs_deflated_" + function_name)
return (methods, attributes)
def _get_bundled_inputs_attributes_and_methods(script_module: torch.jit.ScriptModule) -> Tuple[List[str], List[str]]:
methods: List[str] = []
attributes: List[str] = []
# Has bundled inputs for forward
if hasattr(script_module, 'get_all_bundled_inputs'):
methods.append('get_all_bundled_inputs')
methods.append('get_num_bundled_inputs')
methods.append('run_on_bundled_input')
if hasattr(script_module, 'get_bundled_inputs_functions_and_info'):
methods.append('get_bundled_inputs_functions_and_info')
all_info = script_module.get_bundled_inputs_functions_and_info()
for function_name in all_info:
methods.append("get_all_bundled_inputs_for_" + function_name)
methods.append("_generate_bundled_inputs_for_" + function_name)
attributes.append("_bundled_inputs_deflated_" + function_name)
bundled_inputs_fn = getattr(
script_module,
f"get_all_bundled_inputs_for_{function_name}"
)
num_bundled_inputs: int = len(bundled_inputs_fn())
# Check inflate helper functions for each function, argument and bundled input
func = getattr(script_module, function_name, None)
for arg_idx in range(len(func.schema.arguments) - 1):
for input_idx in range(num_bundled_inputs):
helper_fn_name = _get_inflate_helper_fn_name(
arg_idx=arg_idx,
input_idx=input_idx,
function_name=function_name
)
# if the arg has an InflatableArg with fmt_fn, add the helper function name
if hasattr(script_module, helper_fn_name):
methods.append(helper_fn_name)
return (methods, attributes)