[functorch] update discover_coverage

Author: zou3519

functorch/test/discover_coverage.py

@@ -183,7 +183,7 @@ def safe_append(dct, key, val):
 }
 
 
-def get_top_ops(torch_threshold, nn_fn_threshold):
+def get_top_ops(torch_threshold, nn_fn_threshold, with_counts=False):
     denylist = set({
         # These are either not real "operators", factory functions
         # that trivially work, or not-documented ops.
@@ -228,12 +228,17 @@ def get_top_ops(torch_threshold, nn_fn_threshold):
         'fft',  # is namespace
     })
 
-    torch_ops = [op[0] for op in top_ops.top_torch]
-    nn_fn_ops = [op[0] for op in top_ops.get_nn_functional_top_list()]
-    torch_ops = [op for op in torch_ops if op not in denylist]
-    nn_fn_ops = [op for op in nn_fn_ops if op not in denylist]
+    torch_ops = [op for op in top_ops.top_torch]
+    nn_fn_ops = [op for op in top_ops.get_nn_functional_top_list()]
+    torch_ops = [op for op in torch_ops if op[0] not in denylist]
+    nn_fn_ops = [op for op in nn_fn_ops if op[0] not in denylist]
 
     ops = torch_ops[:torch_threshold] + nn_fn_ops[:nn_fn_threshold]
+
+    # Now, sort by priority
+    ops.sort(reverse=True, key=lambda op: op[1])
+    if not with_counts:
+        ops = [op[0] for op in ops]
     return ops
 
 
@@ -341,8 +346,6 @@ def get_skipped_or_xfailed_ops_for(test_name):
     return result
 
 
-# import pdb; pdb.set_trace()
-
 def get_statuses(for_subset=None, invert=False):
     overridable_outplace_we_care_about = get_public_overridable_outplace_we_care_about()
     if for_subset is not None:
@@ -886,3 +889,8 @@ def summary(self):
 # result = opset.query(Operator.supports_jvpvjp, (Support.NO, Support.UNKNOWN))
 # pprint.pprint(result)
 # print(opset.summary())
+
+# Print list of everything in order
+# all_ops = get_top_ops(999999, 999999, with_counts=True)
+# for op, count in all_ops:
+#     print(f'{op}, {count}')