pytorch
40df3b41 - [AO] Update qLSTM implementation to remove unsupported backend ops (#96436)

Commit
1 year ago
[AO] Update qLSTM implementation to remove unsupported backend ops (#96436) Summary: The reference quantized LSTM implementation uses unbind and inplace squeeze both of which are not supported when building BoltNN's Espresso IR graph. This change adjusts the reference AO Quantizable LSTM implementation without affecting numerically while enabling removal of unsupported ops in BoltNN. Modifications & Adjustments 1. Unbind ops appear when unstacking tensor in loop. Replaced this by getting first dim from shape and looping using ranged index. 2. Removed unbind ops call where the pattern is `[x = t.unbind(0) -> x[i]]` can be just replaced by `t[i]` as creating a tuple from unbind is unnecessary. 3. inplace squeeze `squeeze_` uses which were not required has been replaced by `squeeze`. See notebook N3235193 which was used for testing quantization flow and inspect the torch scripted quantized model for the set of ops used(See last cell). Test Plan: N3235193 Reviewed By: andrewor14 Differential Revision: D43935389 Pull Request resolved: https://github.com/pytorch/pytorch/pull/96436 Approved by: https://github.com/andrewor14
Author
Committer
Parents
Loading