pytorch

f17e02fd - Fix tracing for dynamic quantized LSTM (#29331)

Fix tracing for dynamic quantized LSTM (#29331) Summary: Pull Request resolved: https://github.com/pytorch/pytorch/pull/29331 Closes #27954 This fixes the hard-coding of packed parameter values for the dynamic quantized LSTM by orchestrating the following dance: 1) Each variadic parameter on the module has its own Module. That Module defines the `__getstate__` and __setstate__` method s.t. packed weights are properly re-done on model load. 2) Each of these modules is wrapped into a `torch.nn.ModuleList`, s.t. the parameters appear as attributes in the hierarchy. Then, `gatherParametersAndBuffers` (https://github.com/pytorch/pytorch/blob/9c43b16df9dad3dfb4da1efab68d8c88e6437e8f/torch/csrc/jit/tracer.cpp#L285) can see these parameters and create a `Value*` for them in the traced graph. 3) In forward, we need to convert from ModuleList -> Module -> Parameter to a simple TensorList of the parameters. We just use a loop here. In tracing, we simply record a `ListConstruct` with each of the proper parameter values. In scripting, the `ModuleList` is const, so it can be unrolled into the graph and a subsequent `ListConstruct` does its business. The `forward` of the traced LSTM before and after this change are as follows: Before ``` def forward(self, input: Tensor, argument_2: Tuple[Tensor, Tensor]) -> Tuple[Tensor, Tuple[Tensor, Tensor]]: hx, hx0, = argument_2 _0, _1, _2 = torch.quantized_lstm(input, [hx, hx0], [CONSTANTS.c0, CONSTANTS.c1], True, 1, 0., True, False, False, dtype=12, use_dynamic=True) return (_0, (_1, _2)) ``` After ``` def forward(self, input: Tensor, argument_2: Tuple[Tensor, Tensor]) -> Tuple[Tensor, Tuple[Tensor, Tensor]]: _0 = self.cell._all_weight_values _1 = getattr(_0, "0").param _2 = getattr(_0, "1").param hx, hx0, = argument_2 _3, _4, _5 = torch.quantized_lstm(input, [hx, hx0], [_1, _2], True, 1, 0., True, False, False, dtype=12, use_dynamic=True) return (_3, (_4, _5)) ``` Test Plan: Imported from OSS Differential Revision: D18359880 Pulled By: jamesr66a fbshipit-source-id: 0ff2cad294a1871123015dfc704eaf73a7ac1d9e