Only share the observer for a dropout that is not training

Dropout was added to the list of operators that take their input's observer, and
that list matches on the operator alone. Dropout is only an identity when it is not
training. While training it scales the values it keeps, so its output range is wider
than its input's and a shared observer measures the wrong range. The list also
matches one spelling, so the in-place form was left with the original problem.

Read the training argument off the node, and accept both spellings.

Test Plan:
Confirmed the argument is there to read: exporting a module with the dropout in eval
gives a node whose arguments end in false, and in training gives one ending in true.

Confirmed the in-place form is a distinct operator that survives export, in both
modes, so the previous list genuinely missed it.

Measured the effect on a convolution, dropout, reshape model prepared for
quantization-aware training. With the operator-only list the graph comes out with
five fake-quantize slots collapsed onto three objects, so the training dropout is
sharing. With this change the same graph keeps three slots and three objects, so it
is not.

Checked the predicate over all four combinations: both spellings share when not
training and do not share when training.

Eval-mode sharing is unchanged, which is what the reshape needed, and the operators
already on the list are untouched.

Author: pytorchbot

backends/xnnpack/quantizer/xnnpack_quantizer_utils.py

@@ -1018,7 +1018,7 @@ def _has_annotated_qparam_source(node: Node) -> bool:
         ):
             return True
 
-        if node.op != "call_function" or not _is_share_obs_or_fq_op(node.target):
+        if node.op != "call_function" or not _shares_observer_with_input(node):
             return False
 
         prev_node = node.args[0]
@@ -1105,14 +1105,26 @@ def _is_share_obs_or_fq_op(op: Callable) -> bool:
         torch.ops.aten.slice.Tensor,
         torch.ops.aten.slice_copy.Tensor,
         torch.ops.aten.flatten.using_ints,
-        # Identity in eval mode, so it must not get its own observer.
-        torch.ops.aten.dropout.default,
     ]
 
 
+def _shares_observer_with_input(node: Node) -> bool:
+    """Whether a node must take its input's observer instead of its own.
+
+    Dropout is only an identity when it is not training; while training it scales
+    the values it keeps, so its range is wider than its input's.
+    """
+    if node.target in (
+        torch.ops.aten.dropout.default,
+        torch.ops.aten.dropout_.default,
+    ):
+        return len(node.args) > 2 and node.args[2] is False
+    return _is_share_obs_or_fq_op(node.target)
+
+
 def propagate_annotation(model: torch.fx.GraphModule) -> None:
     for n in model.graph.nodes:
-        if n.op != "call_function" or not _is_share_obs_or_fq_op(n.target):
+        if n.op != "call_function" or not _shares_observer_with_input(n):
             continue
 
         prev_node = n.args[0]