pytorch

0118dec2 - [Pytorch] Expanded Bundled Inputs To Any Public Function (#51153)

[Pytorch] Expanded Bundled Inputs To Any Public Function (#51153) Summary: Pull Request resolved: https://github.com/pytorch/pytorch/pull/51153 Enabled bundled inputs for all public functions that the user wants in a torchscript module. An important caveat here is that you cant add bundled inputs to functions that were in the nn.module but weren't caught in the scripting/tracing process that brought the model to torchscript. Old Api is exactly the same. Still only works on forward, return types the same, etc. -----------New API------------- Attachment of inputs: ***augment_model_with_bundled_inputs*** : works the same as before but added the option to specify an info dictionary. ***augment_many_model_functions_with_bundled_inputs*** : Similar to the above function but allows the user to specify a Dict[Callable, List[<inputs>]] (mapping function references to the bundled inputs for that function) to attach bundled inputs to many functions Consumption of inputs: ***get_all_bundled_inputs_for_<function_name>()*** : Works exactly like get_all_bundled_inputs does, but can be used for functions other then forward if you know ahead of time what they are called, and if they have bundled inputs. ***get_bundled_inputs_functions_and_info()*** : This is easily the hackiest function. Returns a Dict['str', 'str'] mapping function_names to get_all_bundled_inputs_for_<function_name>. A user can then execute the functions specified in the values with something like all_info = model.get_bundled_inputs_functions_and_info() for func_name in all_info.keys(): input_func_name = all_info[func_name]['get_inputs_function_name'][0] func_to_run = getattr(loaded, input_func_name) The reason its done this way is because torchscript doesn't support 'Any' type yet meaning I can't return the bundled inputs directly because they could be different types for each function. Torchscript also doesn't support callable so I can't return a function reference directly either. ghstack-source-id: 120768561 Test Plan: Got a model into torchscript using the available methods that I'm aware of (tracing, scripting, old scripting method). Not really sure how tracing brings in functions that arent in the forward call path though. Attached bundled inputs and info to them successfully. Changes to TorchTest.py on all but the last version of this diff (where it will be/is removed for land) illustrate what I did to test. Created and ran unit test Reviewed By: dreiss Differential Revision: D25931961 fbshipit-source-id: 36e87c9a585554a83a932e4dcf07d1f91a32f046