pytorch/torch/utils
rzou a3cdbd8189 [FlopCounterMode] Fix register_flop_formula (#131777)
Previously, FlopCounterMode would ignore any custom ops registered
through `register_flop_formula`. The problem was:
- register_flop_formula(target) requires target to be an OpOverloadPacket.
- register_flop_formula used register_decomposition to populate its registry
- register_decomposition decomposes the OpOverloadPacket into OpOverload before
  putting it into the registry
- FlopCounterMode ignores OpOverloads in its registry (it assumes the
  registry is a dictionary mapping OpOverloadPacket to flop formula).

register_decomposition is too heavy of a hammer, plus this isn't a
decomposition, so I changed the registration mechanism.

Test Plan:
- new tests
Pull Request resolved: https://github.com/pytorch/pytorch/pull/131777
Approved by: https://github.com/Chillee
2024-07-26 18:44:50 +00:00
..
_strobelight
_sympy
backcompat
benchmark
bottleneck
data
hipify
jit
model_dump
tensorboard
viz
__init__.py
_backport_slots.py backport dataclass(slots=True) (#131014) 2024-07-21 19:26:31 +00:00
_config_module.py Fully type torch/utils/_config_module.py (#131676) 2024-07-24 19:36:09 +00:00
_config_typing.pyi
_content_store.py
_contextlib.py
_cpp_extension_versioner.py
_cxx_pytree.py [pytree] Only import optree if it's used (#131478) 2024-07-24 00:10:49 +00:00
_device.py
_exposed_in.py [BE] typing for decorators - _library/custom_ops (#131578) 2024-07-25 22:24:19 +00:00
_foreach_utils.py
_freeze.py [BE] mypy: disallow untyped decorators (#131428) 2024-07-23 21:50:55 +00:00
_get_clean_triton.py
_import_utils.py
_mode_utils.py
_ordered_set.py Add inductor OrderedSet (#130003) 2024-07-26 18:16:57 +00:00
_python_dispatch.py [BE] typing for decorators - fx/_compatibility (#131568) 2024-07-25 22:24:19 +00:00
_pytree.py [pytree] Only import optree if it's used (#131478) 2024-07-24 00:10:49 +00:00
_stats.py
_traceback.py
_triton.py
_typing_utils.py
_zip.py
backend_registration.py
bundled_inputs.py
checkpoint.py [BE] mypy: disallow untyped decorators (#131428) 2024-07-23 21:50:55 +00:00
collect_env.py
cpp_backtrace.py
cpp_extension.py
deterministic.py
dlpack.py
file_baton.py
flop_counter.py [FlopCounterMode] Fix register_flop_formula (#131777) 2024-07-26 18:44:50 +00:00
hooks.py
mkldnn.py
mobile_optimizer.py
model_zoo.py
module_tracker.py
show_pickle.py
throughput_benchmark.py
weak.py