pytorch

1
# mypy: allow-untyped-defs
2
# NOTE! PLEASE KEEP THIS FILE *FREE* OF TORCH DEPS! IT SHOULD BE IMPORTABLE ANYWHERE.
3
# IF YOU FEEL AN OVERWHELMING URGE TO ADD A TORCH DEP, MAKE A TRAMPOLINE FILE A LA torch._dynamo.utils
4
# AND SCRUB AWAY TORCH NOTIONS THERE.
5
import collections
6
import functools
7
from typing import OrderedDict
8

9
simple_call_counter: OrderedDict[str, int] = collections.OrderedDict()
10

11
def count_label(label):
12
    prev = simple_call_counter.setdefault(label, 0)
13
    simple_call_counter[label] = prev + 1
14

15
def count(fn):
16
    @functools.wraps(fn)
17
    def wrapper(*args, **kwargs):
18
        if fn.__qualname__ not in simple_call_counter:
19
            simple_call_counter[fn.__qualname__] = 0
20
        simple_call_counter[fn.__qualname__] = simple_call_counter[fn.__qualname__] + 1
21
        return fn(*args, **kwargs)
22
    return wrapper
23