Skip to content

GitLab

Commit 6f441f78 authored by Yeqing Li's avatar Yeqing Li Committed by A. Unique TensorFlower
Browse files

Internal change 

PiperOrigin-RevId: 429075761
parent b50cea6e
Hide whitespace changes
Inline Side-by-side
Showing with 11 additions and 2 deletions
official/core/registry.py
View file @ 6f441f78
...@@ -13,6 +13,7 @@ ...@@ -13,6 +13,7 @@
# limitations under the License. # limitations under the License.
"""Registry utility.""" """Registry utility."""
from absl import logging
def register(registered_collection, reg_key): def register(registered_collection, reg_key):
...@@ -54,8 +55,16 @@ def register(registered_collection, reg_key): ...@@ -54,8 +55,16 @@ def register(registered_collection, reg_key):
leaf_reg_key = reg_key leaf_reg_key = reg_key
if leaf_reg_key in collection: if leaf_reg_key in collection:
raise KeyError("Function or class {} registered multiple times.".format( if "beta" in fn_or_cls.__module__:
leaf_reg_key)) # TODO(yeqing): Clean this temporary branch for beta.
logging.warn(
"Duplicate registeration of beta module "
"name %r new %r old %r", reg_key, collection[leaf_reg_key],
fn_or_cls.__module__)
return
else:
raise KeyError("Function or class {} registered multiple times.".format(
leaf_reg_key))
collection[leaf_reg_key] = fn_or_cls collection[leaf_reg_key] = fn_or_cls
return fn_or_cls return fn_or_cls
......
Markdown is supported
0% or .
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment