MCPcopy Create free account
hub / github.com/NVIDIA/TensorRT / register_impl

Method register_impl

tools/Polygraphy/polygraphy/json/serde.py:104–136  ·  view source on GitHub ↗
(func)

Source from the content-addressed store, hash-verified

102 """
103
104 def register_impl(func):
105 def add(key, val):
106 if key in cls.polygraphy_registered:
107 G_LOGGER.critical(
108 f"Duplicate serialization function for type: {key}.\nNote: Existing function: {cls.polygraphy_registered[key]}, New function: {func}"
109 )
110 cls.polygraphy_registered[key] = val
111
112 if cls == Encoder:
113
114 def wrapped(obj):
115 dct = func(obj)
116 dct[constants.TYPE_MARKER] = str_from_type(typ)
117 return dct
118
119 add(typ, wrapped)
120 return wrapped
121 elif cls == Decoder:
122
123 def wrapped(dct):
124 if constants.TYPE_MARKER in dct:
125 del dct[constants.TYPE_MARKER]
126
127 type_name = legacy_str_from_type(typ)
128 if type_name in dct:
129 del dct[type_name]
130
131 return func(dct)
132
133 add(legacy_str_from_type(typ), wrapped)
134 add(str_from_type(typ), wrapped)
135 else:
136 G_LOGGER.critical("Cannot register for unrecognized class type: ")
137
138 return register_impl
139

Callers

nothing calls this directly

Calls 4

legacy_str_from_typeFunction · 0.85
str_from_typeFunction · 0.85
addFunction · 0.50
criticalMethod · 0.45

Tested by

no test coverage detected