[FX] Fix operator_schemas normalize_function to consider OpOverloads (p…

…ytorch#76469)

Summary:
Pull Request resolved: pytorch#76469

Broken by Original commit changeset: 450e86c4e08a

Original Phabricator Diff: D35874477

Test Plan: Added unit test coverage to test_fx_experimental

Reviewed By: albanD

Differential Revision: D35978105

fbshipit-source-id: f22670b3b00a86777a26feaf4cb911595d150a17
(cherry picked from commit 91868b1)

test/test_fx_experimental.py

@@ -1642,6 +1642,33 @@ def forward(self, {', '.join(param_names)}):
             test_out = traced(*param_values)
             self.assertEqual(test_out, ref_out)
 
+    def test_normalize_quantized_eb(self):
+        target = torch.ops.quantized.embedding_bag_byte_rowwise_offsets
+        args = (
+            torch.empty((2, 3), dtype=torch.uint8),
+            torch.empty((2,), dtype=torch.int64),
+            torch.empty((2,), dtype=torch.int64),
+        )
+        norm_args_and_kwargs = normalize_function(
+            target, args, normalize_to_only_use_kwargs=True
+        )
+        self.assertTrue(norm_args_and_kwargs is not None)
+        self.assertEqual(
+            set(norm_args_and_kwargs.kwargs.keys()),
+            {
+                "weight",
+                "indices",
+                "offsets",
+                "scale_grad_by_freq",
+                "mode",
+                "pruned_weights",
+                "per_sample_weights",
+                "compressed_indices_mapping",
+                "include_last_offset",
+            },
+        )
+        self.assertEqual(norm_args_and_kwargs.args, tuple())
+
 
 instantiate_device_type_tests(TestNormalizeOperators, globals())
 

torch/fx/operator_schemas.py

@@ -257,6 +257,8 @@ def normalize_function(
     if kwargs is None:
         kwargs = {}
     new_args_and_kwargs = None
+    if isinstance(target, OpOverloadPacket) or isinstance(target, OpOverload):
+        target = target.op
     if not isinstance(target, types.BuiltinFunctionType):
         target_for_analysis = target
         if target in boolean_dispatched: