Generalize QuantizedOutputWrapper for multi-output models (#19987)

Summary:
Pull Request resolved: #19987

Add support for multiple outputs in quantized output wrapper

Differential Revision: D107429509

Author: ethansfng

backends/cadence/aot/compiler_funcs.py

@@ -400,33 +400,55 @@ def sink_dequants(program: torch.export.ExportedProgram) -> None:
 
 class QuantizedOutputWrapper(torch.nn.Module):
     """
-    Wrapper that quantizes a model's output so it produces uint8 tensors.
+    Wrapper that quantizes a model's output(s) so they produce quantized tensors.
 
     Mirrors QuantizedInputWrapper: the wrapper adds a quantize_per_tensor after
-    the model's output. When the graph is traced, the dequant (from the model) →
+    each output. When the graph is traced, the dequant (from the model) →
     quant (from the wrapper) pair with matching parameters folds away, leaving
     the output in its quantized form.
 
     Args:
         module: The module to wrap (may already be a QuantizedInputWrapper).
-        output_quant_args: (scale, zero_point, qmin, qmax, dtype) for the output.
+        output_quant_args: Quantization parameters — either a single QuantArgs
+            tuple or a list with one entry per output.
     """
 
     def __init__(
         self,
         module: torch.nn.Module,
-        output_quant_args: QuantArgs,
+        output_quant_args: Union[QuantArgs, list[QuantArgs | None]],
     ) -> None:
         super().__init__()
         self.module: torch.nn.Module = module
-        self.output_quant_args: QuantArgs = output_quant_args
+        if isinstance(output_quant_args, list):
+            self._multi_output: bool = True
+            self._per_output_args: list[QuantArgs | None] = output_quant_args
+        else:
+            self._multi_output = False
+            self._per_output_args = [output_quant_args]
 
     def forward(self, *args: torch.Tensor) -> Any:
         result = self.module(*args)
-        scale, zp, qmin, qmax, dtype = self.output_quant_args
-        return torch.ops.quantized_decomposed.quantize_per_tensor.default(
-            result, scale, zp, qmin, qmax, dtype
-        )
+        if not self._multi_output:
+            qa = self._per_output_args[0]
+            assert qa is not None
+            scale, zp, qmin, qmax, dtype = qa
+            return torch.ops.quantized_decomposed.quantize_per_tensor.default(
+                result, scale, zp, qmin, qmax, dtype
+            )
+        out: list[torch.Tensor] = []
+        for i, r in enumerate(result):
+            qa = self._per_output_args[i] if i < len(self._per_output_args) else None
+            if qa is None:
+                out.append(r)
+                continue
+            scale, zp, qmin, qmax, dtype = qa
+            out.append(
+                torch.ops.quantized_decomposed.quantize_per_tensor.default(
+                    r, scale, zp, qmin, qmax, dtype
+                )
+            )
+        return tuple(out)
 
 
 def _get_transparent_ops() -> set[Any]: