import torch import copy from torch.testing._internal.common_methods_invocations import op_db from functorch_additional_op_db import additional_op_db from enum import Enum import functorch._src.top_operators_github_usage as top_ops import pprint import unittest import enum from torch.testing._internal.common_device_type import toleranceOverride # Importing these files make modifications to the op_db that we need import test_ops # noqa: F401 import test_vmap # noqa: F401 all_overridable = list(torch.overrides.get_testing_overrides().keys()) public_docs = [ (torch.nn.functional, 'torch.nn.functional', 'docs/source/nn.functional.rst'), (torch.fft, 'torch.fft', 'docs/source/fft.rst'), (torch.special, 'torch.special', 'docs/source/special.rst'), (torch.linalg, 'torch.linalg', 'docs/source/linalg.rst'), (torch, 'torch', 'docs/source/torch.rst'), (torch.Tensor, 'torch.Tensor', 'docs/source/tensors.rst'), ] # torch.abs, Tensor.abs, Tensor.abs_ are all considered to be different def get_public_overridable_apis(pytorch_root='/raid/rzou/pt/debug-cpu'): results = {} all_overridable_apis = set(torch.overrides.get_testing_overrides().keys()) for module, module_name, src in public_docs: with open(f'{pytorch_root}/{src}') as f: lines = f.readlines() # APIs eitehr begin with 4 spaces or ".. autofunction::" api_lines1 = [line.strip() for line in lines if line.startswith(' ' * 4)] api_lines2 = [line.strip()[len('.. autofunction:: '):] for line in lines if line.startswith('.. autofunction::')] lines = api_lines1 + api_lines2 lines = [line[7:] if line.startswith('Tensor.') else line for line in lines] lines = [line for line in lines if hasattr(module, line)] for line in lines: api = getattr(module, line) if api in all_overridable_apis: results[f'{module_name}.{line}'] = api return results denylist = { 'torch.Tensor.data_ptr', 'torch.Tensor.dim', 'torch.Tensor.element_size', 'torch.Tensor.backward', 'torch.Tensor.as_strided', 'torch.Tensor.register_hook', 'torch.Tensor.record_stream', 'torch.Tensor.qscheme', 'torch.Tensor.ndimension', 'torch.Tensor.smm', 'torch.Tensor.sspaddmm', 'torch.Tensor.retain_grad', 'torch.Tensor.sparse_mask', 'torch.Tensor.sparse_dim', 'torch.Tensor.dense_dim', 'torch.Tensor.values', 'torch.Tensor.indices', 'torch.Tensor.numel', 'torch.Tensor.size', 'torch.Tensor.nelement', 'torch.Tensor.q_scale', 'torch.Tensor.q_zero_point', 'torch.Tensor.q_per_channel_scales', 'torch.Tensor.q_per_channel_zero_points', 'torch.Tensor.q_per_channel_axis', 'torch.Tensor.int_repr', 'torch.Tensor.to_sparse', 'torch.Tensor.is_inference', 'torch.Tensor.storage', 'torch.Tensor.storage_type', } def get_method_only_ops_we_care_about(): apis = get_public_overridable_apis() result = [] for key, _ in apis.items(): if not key.startswith('torch.Tensor'): continue if key in denylist: continue api = key.split('.')[2] # filter out in-place if api.endswith('_'): continue if f'torch.{api}' not in apis.keys(): result.append(api) return result # Deduplicates torch.abs and Tensor.abs def get_public_overridable_ops(): results = get_public_overridable_apis() cpy = copy.deepcopy(results) for key, _ in cpy.items(): if not key.startswith('torch.Tensor'): continue api = key.split('.')[2] if f'torch.{api}' in results.keys(): del results[key] return results def get_public_overridable_outplace_ops(): results = get_public_overridable_ops() cpy = copy.deepcopy(results) for key, _ in cpy.items(): # NB: there are no dunder methods bcs we don't document those if key.endswith('_'): del results[key] return results def get_public_overridable_outplace_we_care_about(): results = get_public_overridable_outplace_ops() cpy = copy.deepcopy(results) for key, _ in cpy.items(): # quantization if 'quant' in key or '.q_' in key: del results[key] # is_cpu, etc. It doesn't make sense to have OpInfos for these if '.is_' in key: del results[key] if key in denylist and key in results: del results[key] return results # e.g. nn.functional.softmax def get_op(dotted_name): names = dotted_name.split('.') mod = torch for name in names: if not hasattr(mod, name): return None mod = getattr(mod, name) return mod # Maps function -> [OpInfo] def get_ops_covered_by_opinfos(): ops = {} def safe_append(dct, key, val): if key in dct: dct[key].append(val) else: dct[key] = [val] for opinfo in op_db: func_op = get_op(opinfo.name) if func_op: safe_append(ops, func_op, opinfo) if opinfo.method_variant: safe_append(ops, opinfo.method_variant, opinfo) if opinfo.inplace_variant: safe_append(ops, opinfo.inplace_variant, opinfo) for alias in opinfo.aliases: safe_append(ops, alias.op, opinfo) return ops factory_fns = { 'tensor', 'zeros', 'ones', 'randn', 'arange', 'rand', 'empty', 'randperm', 'linspace', 'logspace', 'hann_window', 'full', 'eye', 'blackman_window', 'barlett_window', 'randint', 'range', 'arange', } 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. 'load', 'no_grad', 'save', 'from_numpy', 'manual_seed', 'set_grad_enabled', 'set_default_tensor_type', 'set_num_threads', 'set_printoptions', 'numel', 'set_default_dtype', 'sparse_coo_tensor', 'set_rng_state', 'get_rng_state', 'get_default_dtype', 'initial_seed', 'get_num_threads', 'quantize_per_tensor', 'hann_window', 'is_tensor', 'as_tensor', 'equal', 'enable_grad', 'seed', 'is_storage', 'is_floating_point', 'nn.functional.torch', 'set_flush_denormal', 'set_num_interop_threads', 'dequantize', 'get_num_interop_threads', 'nn.functional.math', 'nn.functional.threshold_', 'nn.functional.selu_', 'nn.functional.elu_', 'nn.functional.rrelu_', 'nn.functional.leaky_relu_', 'nn.functional.hardtanh_', 'nn.functional.has_torch_function', 'nn.functional.has_torch_function_unary', 'nn.functional.has_torch_function_variadic', 'nn.functional.handle_torch_function', 'nn.functional.adaptive_max_pool1d_with_indices', 'nn.functional.adaptive_max_pool2d_with_indices', 'nn.functional.adaptive_max_pool3d_with_indices', 'nn.functional.fractional_max_pool2d_with_indices', 'nn.functional.fractional_max_pool3d_with_indices', 'is_complex', 'grad', 'quantize_per_channel', 'nn.functional.max_pool2d_with_indices', 'nn.functional.max_pool3d_with_indices', 'nn.functional.max_pool1d_with_indices', 'nn.functional.celu_', 'nn.functional.grad', 'nn.functional.relu_', 'nn.functional.boolean_dispatch', 'nn.functional.assert_int_or_pair', 'fft', # is namespace }) torch_ops = top_ops.top_torch nn_fn_ops = 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 def get_ops_percentage(torch_threshold, nn_fn_threshold): data = top_ops.top_torch + top_ops.get_nn_functional_top_list() def get_num_usages(opname): # Ignore this, this is heavily inflated if opname == 't': return 0 result = [op[1] for op in data if op[0] == opname] assert len(result) == 1 return result[0] # get all operators that are not in the denylist all_ops = get_top_ops(999999, 999999) total_op_usages = sum([get_num_usages(op) for op in all_ops]) # get subset of all operators subset_ops = get_top_ops(torch_threshold, nn_fn_threshold) subset_op_usages = sum([get_num_usages(op) for op in subset_ops]) return subset_op_usages / total_op_usages def get_top_ops_not_covered_by_opinfo(torch_threshold=0, nn_fn_threshold=0): ops = get_top_ops(torch_threshold, nn_fn_threshold) ops_with_opinfo = [] for op in op_db: ops_with_opinfo.append(op.name) ops_with_opinfo.extend([op.name for op in op.aliases]) ops_with_opinfo = set(ops_with_opinfo) result = [op for op in ops if op not in ops_with_opinfo] result = [op for op in result if op not in denylist] result = [op for op in result if op not in factory_fns] return result def get_covered_ops(ops_list, invert=False): ops_covered_by_opinfo = get_ops_covered_by_opinfos() overridable_outplace_ops = ops_list results = {} for key, op in overridable_outplace_ops.items(): cond = op in ops_covered_by_opinfo if invert: cond = not cond if cond: results[key] = op return results class Status(Enum): Correct = 0 Fast = 1 tests = { 'test_vmap_exhaustive', 'test_op_has_batch_rule', 'test_vjp', 'test_vmapvjp', 'test_vmapvjp_has_batch_rule', 'test_jvp', 'test_vmapjvp', } def is_decorateinfo_skip_or_xfail(decorateinfo): assert len(decorateinfo.decorators) == 1 actual_decorator = decorateinfo.decorators[0] if isinstance(actual_decorator, toleranceOverride): return False if actual_decorator == unittest.expectedFailure: return True # Assume the rest are skips return True def get_all_tested_ops(): overridable_outplace_we_care_about = get_public_overridable_outplace_we_care_about() op_to_opinfo = get_ops_covered_by_opinfos() result = set({}) for name, op in get_covered_ops(overridable_outplace_we_care_about).items(): opinfos = op_to_opinfo[op] for opinfo in opinfos: result.add(opinfo.name) return result def get_skipped_or_xfailed_ops_for(test_name): overridable_outplace_we_care_about = get_public_overridable_outplace_we_care_about() op_to_opinfo = get_ops_covered_by_opinfos() result = set({}) for name, op in get_covered_ops(overridable_outplace_we_care_about).items(): opinfos = op_to_opinfo[op] for opinfo in opinfos: for decorator in opinfo.decorators: if not hasattr(decorator, 'test_name'): continue if decorator.test_name != test_name: continue if is_decorateinfo_skip_or_xfail(decorator): result.add(opinfo.name) return result 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: overridable_outplace_we_care_about = { k: v for k, v in overridable_outplace_we_care_about.items() # Removes "torch." if k[6:] in for_subset } op_to_opinfo = get_ops_covered_by_opinfos() result = {} _ = get_covered_ops(overridable_outplace_we_care_about) def get_covered_tests(op): opinfos = op_to_opinfo[op] result = copy.deepcopy(tests) for opinfo in opinfos: for decorator in opinfo.decorators: if not hasattr(decorator, 'test_name'): continue if decorator.test_name in tests and decorator.test_name in result: result.remove(decorator.test_name) return result def get_all_aliases(op): opinfos = op_to_opinfo[op] result = [] for opinfo in opinfos: result.append(opinfo.name) result.extend(opinfo.aliases) return set(result) for name, op in get_covered_ops(overridable_outplace_we_care_about).items(): successful_tests = get_covered_tests(op) failed_tests = tests - successful_tests result[name] = failed_tests if invert else successful_tests return result def transpose_statuses(for_subset=None, invert=False): statuses = get_statuses(for_subset, invert=invert) result = {} for test in tests: result[test] = set({}) for op, supported in statuses.items(): for test in supported: result[test].add(op) return result overridable_apis = get_public_overridable_apis() overridable_ops = get_public_overridable_ops() overridable_outplace_ops = get_public_overridable_outplace_ops() overridable_outplace_we_care_about = get_public_overridable_outplace_we_care_about() tested_overridable_outplace_ops = get_covered_ops(overridable_outplace_we_care_about) untested_overridable_outplace_ops = get_covered_ops(overridable_outplace_we_care_about, invert=True) # print("List of OpInfos we need:") # for key in untested_overridable_outplace_ops.keys(): # print(key) # print("-" * 80) # print("") print(f'Overridable public APIs: {len(overridable_apis)}') print(f'Overridable public ops: {len(overridable_ops)}') print(f'Overridable public outplace ops: {len(overridable_outplace_ops)}') print(f'Overridable public outplace ops we care about: {len(overridable_outplace_we_care_about)}') print(f'OpInfo-tested overridable public outplace ops: {len(tested_overridable_outplace_ops)}') def remove_torch(name): assert name[:6] == 'torch.' return name[6:] def get_list_of_all_tests(): all_tests = list(tested_overridable_outplace_ops.keys()) return set([remove_torch(test) for test in all_tests]) mytest = { 'test_vmap_exhaustive', 'test_op_has_batch_rule', 'test_vjp', 'test_vmapvjp', 'test_vmapvjp_has_batch_rule', } print('*' * 80) all_tests = get_list_of_all_tests() for test in mytest: result = get_skipped_or_xfailed_ops_for(test) diff = len(all_tests - result) print(f'{test}: {diff}') def get_jvp_coverage(subset=None): # - number that support autograd # - number that support forward_ad (in pytorch core) # - number that support functorch.jvp op_to_opinfo = get_ops_covered_by_opinfos() ops_dct = tested_overridable_outplace_ops if subset is not None: ops_dct = {name: op for name, op in ops_dct.items() if remove_torch(name) in subset} supports_autograd_ops_dct = {name: op_to_opinfo[fn] for name, fn in ops_dct.items() if op_to_opinfo[fn][0].supports_autograd} supports_forwardad_ops_dct = {name: op_to_opinfo[fn] for name, fn in ops_dct.items() if op_to_opinfo[fn][0].supports_forward_ad} ops = set([remove_torch(test) for test in list(ops_dct.keys())]) supports_autograd = set([remove_torch(test) for test in list(supports_autograd_ops_dct.keys())]) supports_forward_ad = set([remove_torch(test) for test in list(supports_forwardad_ops_dct.keys())]) assert supports_forward_ad.issubset(supports_autograd) assert supports_autograd.issubset(ops) failed_ops = get_skipped_or_xfailed_ops_for('test_jvp') coverage = len(supports_forward_ad - failed_ops) no_forward_ad = len(supports_autograd) - len(supports_forward_ad) print(f'test_jvp, {coverage}, {no_forward_ad}, {len(ops)}') get_jvp_coverage() get_jvp_coverage(get_top_ops(100, 25)) for op in get_top_ops(100, 25): print(op) print('*' * 80) # result = get_skipped_or_xfailed_ops_for('test_vmap_exhaustive') # result = get_skipped_or_xfailed_ops_for('test_op_has_batch_rule') # result = get_skipped_or_xfailed_ops_for('test_vjp') # result = get_skipped_or_xfailed_ops_for('test_vmapvjp') # result = get_skipped_or_xfailed_ops_for('test_vmapvjp_has_batch_rule') # import pdb; pdb.set_trace() statuses = transpose_statuses() for test in tests: print(f'{test} coverage {len(statuses[test])}') method_only_ops = get_method_only_ops_we_care_about() # for op in method_only_ops: # print(f' {op},') top_ops_not_covered_by_opinfo = get_top_ops_not_covered_by_opinfo(100, 25) print('=' * 80) for op in top_ops_not_covered_by_opinfo: print(f'{op}, {top_ops.usage_count[op]}') # print("top ops not covered by opinfo: ") # top_ops_not_covered_by_opinfo = get_top_ops_not_covered_by_opinfo(200, 50) # for op in top_ops_not_covered_by_opinfo: # print(f'{op}, {top_ops.usage_count[op]}') # print("top ops not covered by opinfo: ") # top_ops_not_covered_by_opinfo = get_top_ops_not_covered_by_opinfo(220, 92) # for op in top_ops_not_covered_by_opinfo: # print(f'{op}, {top_ops.usage_count[op]}') # print("top ops not covered by opinfo: ") # top_ops_not_covered_by_opinfo = get_top_ops_not_covered_by_opinfo(999, 999) # for op in top_ops_not_covered_by_opinfo: # print(f'{op}, {top_ops.usage_count[op]}') def remove_from_set(parent, to_remove): for to_remove_elt in to_remove: if to_remove_elt in parent: parent.remove(to_remove_elt) def print_coverage_info(th=100, nn=25): print('=' * 80) print(f"top {th}, {nn} coverage") statuses = transpose_statuses(get_top_ops(th, nn), invert=True) top_ops_not_covered_by_opinfo = get_top_ops_not_covered_by_opinfo(th, nn) # testing problems exemptions = { 'torch.nn.functional.dropout', # randomness } # Allowed exemptions vmap_exemptions = { 'torch.randn_like', # randomness 'torch.rand_like', # randomness 'torch.allclose', # number output 'torch.unique', # dynamic 'torch.nonzero', # dynamic 'torch.masked_select', # dynamic 'torch.prod', # dynamic (backward) 'torch.norm', # norm with nuc is not commonly used; we support the other cases. 'torch.svd', # There isn't a bug, it is just nondeterministic so we can't test it. 'torch.nn.functional.embedding', # We support everything except the sparse option. } remove_from_set(statuses['test_vmap_exhaustive'], vmap_exemptions) remove_from_set(statuses['test_vmapvjp'], vmap_exemptions) remove_from_set(statuses['test_vmapvjp_has_batch_rule'], vmap_exemptions) remove_from_set(statuses['test_op_has_batch_rule'], vmap_exemptions) remove_from_set(statuses['test_vmapjvp'], vmap_exemptions) for test in tests: remove_from_set(statuses[test], exemptions) print(f"total ops in set: {th + nn}") print(f"tested by OpInfo: {th + nn - len(top_ops_not_covered_by_opinfo)}") for test in tests: if test in {'test_jvp', 'test_vmapjvp'}: continue print(f'{test} failing coverage {len(statuses[test])}') # We don't care about these yet del statuses['test_jvp'] del statuses['test_vmapjvp'] pprint.pprint(statuses) def get_name_to_opinfo_map(): dct = {} for op in (op_db + additional_op_db): def add(name, op): if name not in dct: dct[name] = [] dct[name].append(op) add(op.name, op) for alias in op.aliases: add(alias.name, op) return dct NAME_TO_OPINFO = get_name_to_opinfo_map() class Support(enum.Enum): NO = 0 YES = 1 UNKNOWN = 2 FACTORY_FNS = { 'tensor', 'zeros', 'ones', 'randn', 'arange', 'rand', 'empty', 'range', 'full', 'randperm', 'eye', 'randint', 'linspace', 'logspace', } VJP_EXEMPTIONS = { 'nn.functional.dropout', # not actually problem, randomness testing artifact 'nn.functional.dropout2d', # not actually problem, randomness testing artifact 'nn.functional.rrelu', # not actually problem, randomness testing artifact 'bernoulli', # not actually problem, randomness testing artifact 'normal', # not actually problem, randomness testing artifact } VMAP_EXEMPTIONS = { 'randn_like', # randomness 'rand_like', # randomness 'allclose', # number output 'unique', # dynamic 'nonzero', # dynamic 'masked_select', # dynamic 'prod', # dynamic (backward) 'norm', # norm with nuc is not commonly used; we support the other cases. 'svd', # There isn't a bug, it is just nondeterministic so we can't test it. 'nn.functional.embedding', # We support everything except the sparse option. 'nn.functional.dropout', # randomness 'nn.functional.dropout2d', # randomness 'bernoulli', # randomness 'multinomial', # randomness 'normal', # randomness } JVP_EXEMPTIONS = { 'nn.functional.dropout', # not actually problem, randomness testing artifact 'nn.functional.dropout2d', # not actually problem, randomness testing artifact 'nn.functional.rrelu', # not actually problem, randomness testing artifact 'normal', # not actually problem, randomness testing artifact 'bernoulli', # not actually problem, randomness testing artifact } class Operator: def __init__(self, name): self.name = name self.opinfos = NAME_TO_OPINFO.get(name, None) assert self.opinfos is None or len(self.opinfos) > 0 def has_opinfo(self): return self.opinfos is not None def __repr__(self): return f'Operator("{self.name}")' def __hash__(self): return hash(self.name) def no_opinfos_skip_test(self, test_name): """Returns NO if any opinfos have a skip or xfail for the test""" if not self.has_opinfo(): return Support.UNKNOWN for opinfo in self.opinfos: for decorator in opinfo.decorators: if not hasattr(decorator, 'test_name'): continue if decorator.test_name != test_name: continue if is_decorateinfo_skip_or_xfail(decorator): return Support.NO return Support.YES def any_opinfo_attr(self, attr): if not self.has_opinfo(): raise RuntimeError() return any([getattr(opinfo, attr) for opinfo in self.opinfos]) def all_opinfo_attr(self, attr): if not self.has_opinfo(): raise RuntimeError() return all([getattr(opinfo, attr) for opinfo in self.opinfos]) def supports_vjp(self): if self.name in FACTORY_FNS: return Support.YES if self.name in VJP_EXEMPTIONS: return Support.YES return self.no_opinfos_skip_test('test_vjp') def supports_vmap(self): if self.name in FACTORY_FNS: return Support.YES if self.name in VMAP_EXEMPTIONS: return Support.YES return self.no_opinfos_skip_test('test_vmap_exhaustive') def supports_fast_vmap(self): if self.name in FACTORY_FNS: return Support.YES if self.name in VMAP_EXEMPTIONS: return Support.YES return self.no_opinfos_skip_test('test_op_has_batch_rule') def supports_vmapvjp(self): if self.name in FACTORY_FNS: return Support.YES if self.name in VMAP_EXEMPTIONS: return Support.YES return self.no_opinfos_skip_test('test_vmapvjp') def supports_fast_vmapvjp(self): if self.name in FACTORY_FNS: return Support.YES if self.name in VMAP_EXEMPTIONS: return Support.YES return self.no_opinfos_skip_test('test_vmapvjp_has_batch_rule') def supports_jvp(self): if self.name in FACTORY_FNS: return Support.YES if self.name in JVP_EXEMPTIONS: return Support.YES if not self.has_opinfo(): return Support.UNKNOWN if self.any_opinfo_attr('supports_autograd') and \ not self.all_opinfo_attr('supports_forward_ad'): return Support.NO return self.no_opinfos_skip_test('test_jvp') def supports_jvpvjp(self): if self.name in FACTORY_FNS: return Support.YES exemptions = { # we have support (see OpInfo), testing artifact 'nn.functional.dropout2d', 'nn.functional.dropout', # exception: we dont even support double backward for this 'nn.functional.hardswish', 'bernoulli', # this isn't differentiable 'normal', # not differentiable } if self.name in exemptions: return Support.YES return self.no_opinfos_skip_test('test_jvpvjp') def _supports_vmapjvp_base(self, test): if self.name in FACTORY_FNS: return Support.YES VMAPJVP_EXEMPTIONS = { 'prod', # dynamic (backward) 'nn.functional.batch_norm', # testing problem 'normal', # not actually problem, randomness testing artifact 'bernoulli', # not actually problem, randomness testing artifact 'nn.functional.dropout2d', # not actually problem, randomness testing artifact 'nn.functional.dropout', # not actually problem, randomness testing artifact # Not a problem. # It's just that the max_norm testing mutates inputs... # (we have our own functorch variant of the OpInfo without max_norm) 'nn.functional.embedding', } if self.name in VMAPJVP_EXEMPTIONS: return Support.YES if not self.has_opinfo(): return Support.UNKNOWN if self.any_opinfo_attr('supports_autograd') and \ not self.all_opinfo_attr('supports_forward_ad'): return Support.NO return self.no_opinfos_skip_test(test) def supports_vmapjvp(self): return self._supports_vmapjvp_base('test_vmapjvpall') def supports_fast_vmapjvp(self): return self._supports_vmapjvp_base('test_vmapjvpall_has_batch_rule') class OperatorSet: def __init__(self, operators): self.data = set(operators) @classmethod def from_names(cls, names): return OperatorSet([Operator(name) for name in names]) @classmethod def from_top_ops_threshold(cls, torch_threshold, nn_fn_threshold): names = get_top_ops(torch_threshold, nn_fn_threshold) return cls.from_names(names) @classmethod def from_top125(cls): return cls.from_top_ops_threshold(100, 25) @classmethod def from_top160(cls): return cls.from_top_ops_threshold(107, 53) @classmethod def all(cls): dct = get_public_overridable_outplace_we_care_about() names = dct.keys() names_sanitized = [] for n in names: torch_tensor = 'torch.Tensor.' torch_dot = 'torch.' if n.startswith(torch_tensor): names_sanitized.append(n[len(torch_tensor):]) elif n.startswith(torch_dot): names_sanitized.append(n[len(torch_dot):]) else: raise AssertionError() return cls.from_names(names_sanitized) def query(self, operator_method, filter=(Support.NO, Support.YES, Support.UNKNOWN)): result = {} for key in filter: result[key] = set([]) for op in self.data: support_status = operator_method(op) if support_status in filter: result[support_status].add(op) return result def summary(self): checks = [ 'supports_vjp', 'supports_vmap', 'supports_fast_vmap', 'supports_vmapvjp', 'supports_fast_vmapvjp', 'supports_jvp', 'supports_vmapjvp', 'supports_fast_vmapjvp', 'supports_jvpvjp', ] result = ['test, yes, no, unknown'] for check in checks: accessor = getattr(Operator, check) all_results = self.query(accessor) yes_amt = len(all_results[Support.YES]) no_amt = len(all_results[Support.NO]) unknown_amt = len(all_results[Support.UNKNOWN]) result.append(f'{check}, {yes_amt}, {no_amt}, {unknown_amt}') return '\n'.join(result) opset = OperatorSet.all() has_no_opinfo = opset.query(Operator.has_opinfo, (False,)) print("=" * 30 + " Summary " + "=" * 30) print(f'% of usages on github: {get_ops_percentage(99999, 99999)}') print(opset.summary()) # sanity checks result = opset.query(Operator.supports_vjp, (Support.NO, Support.UNKNOWN)) # pprint.pprint(result) print("=" * 30 + " Top 60 Summary " + "=" * 30) print(f'% of usages on github: {get_ops_percentage(35, 25)}') opset = OperatorSet.from_top_ops_threshold(35, 25) # result = opset.query(Operator.supports_vmapjvp, (Support.NO, Support.UNKNOWN)) # pprint.pprint(result) # result = opset.query(Operator.supports_jvp, (Support.NO, Support.UNKNOWN)) # pprint.pprint(result) # kresult = opset.query(Operator.supports_jvpvjp, (Support.NO, Support.UNKNOWN)) # kpprint.pprint(result) # result = opset.query(Operator.supports_vmapjvp, (Support.NO, Support.UNKNOWN)) # pprint.pprint(result) # result = opset.query(Operator.supports_fast_vmapjvp, (Support.NO, Support.UNKNOWN)) # pprint.pprint(result) # pprint.pprint(result) print(opset.summary()) print("=" * 30 + " Top 125 Summary " + "=" * 30) print(f'% of usages on github: {get_ops_percentage(100, 25)}') opset = OperatorSet.from_top125() # result = opset.query(Operator.supports_vmap, (Support.NO, Support.UNKNOWN)) # pprint.pprint(result) # result = opset.query(Operator.supports_jvpvjp, (Support.NO, Support.UNKNOWN)) # pprint.pprint(result) print("supports_vjp") result = opset.query(Operator.supports_vjp, (Support.NO, Support.UNKNOWN)) pprint.pprint(result) print("supports_jvp") result = opset.query(Operator.supports_jvp, (Support.NO, Support.UNKNOWN)) pprint.pprint(result) print("supports_vmapjvp") result = opset.query(Operator.supports_vmapjvp, (Support.NO, Support.UNKNOWN)) pprint.pprint(result) print("supports_jvpvjp") result = opset.query(Operator.supports_jvpvjp, (Support.NO, Support.UNKNOWN)) pprint.pprint(result) # result = opset.query(Operator.supports_fast_vmapjvp, (Support.NO, Support.UNKNOWN)) # pprint.pprint(result) # pprint.pprint(result) print(opset.summary()) # print("=" * 30 + " Top 160 Summary " + "=" * 30) # opset = OperatorSet.from_top160() # 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}')