/
githubmirror
/
ColossalAI
Обзор
Документация
Войти
/
githubmirror
/
ColossalAI
Код
Запросы
0
Пакеты
0
Релизы
0
Аналитика
Безопасность
main
colossalai/auto_parallel/tensor_shard/node_handler/registry.py
27 строк
744 B
Hongxin Liu
[misc] update pre-commit and run all files (#4752)
19 сен 2023, 09:20
Не верифицирован
19 сен 2023, 09:20
079bf3c
Код
Авторство
О чём код?
class Registry: def __init__(self, name): self.name = name self.store = {} def register(self, source): def wrapper(func): if isinstance(source, (list, tuple)): # support register a list of items for this func for element in source: self.store[element] = func else: self.store[source] = func return func return wrapper def get(self, source): assert source in self.store, f"{source} not found in the {self.name} registry" target = self.store[source] return target def has(self, source): return source in self.store operator_registry = Registry("operator")