diff --git a/quax/_core.py b/quax/_core.py index 2b6c8d5..215c3cd 100644 --- a/quax/_core.py +++ b/quax/_core.py @@ -55,10 +55,13 @@ def _register(rule: Callable): existing_rule = _rules[primitive] # pyright: ignore except KeyError: - @plum.Dispatcher().abstract def existing_rule(): assert False + existing_rule.__name__ = f"{primitive}_dispatcher" + existing_rule.__qualname__ = f"{primitive}_dispatcher" + existing_rule = plum.Dispatcher().abstract(existing_rule) + _rules[primitive] = existing_rule existing_rule.dispatch(rule) return rule