[PyTorch] Use torch's register_custom_class API for opaque quantizers when available

tests/pytorch/test_torch_compile.py

@@ -22,7 +22,7 @@
     from torch._opaque_base import OpaqueBaseMeta
     from torch._library.opaque_object import (
         get_opaque_type_name,
-        register_opaque_type,
+        register_custom_class,
         MemberType,
     )
 
@@ -215,9 +215,9 @@ def __fx_repr__(self):
                 {"ToyQuantizer": ToyQuantizer},
             )
 
-    register_opaque_type(
+    register_custom_class(
         ToyQuantizer,
-        typ="value",
+        typ="constant",
         members={
             "__setattr__": MemberType.USE_REAL,
             "set_usage": MemberType.USE_REAL,

transformer_engine/pytorch/dynamo/custom_op.py

@@ -180,7 +180,7 @@ def is_simple_value(cls, value: Any) -> bool:
             return True
         if isinstance(value, type):
             return True
-        if _is_opaque_value_type is not None and _is_opaque_value_type(type(value)):
+        if _is_opaque_constant_type is not None and _is_opaque_constant_type(type(value)):
             return True
         if isinstance(value, dict):
             return all(isinstance(k, str) and cls.is_simple_value(v) for k, v in value.items())
@@ -223,7 +223,7 @@ def _fmt_simple(cls, value: Any) -> str:
         if isinstance(value, tuple):
             body = ", ".join(cls._fmt_simple(v) for v in value)
             return f"({body},)" if len(value) == 1 else f"({body})"
-        if _is_opaque_value_type(type(value)):
+        if _is_opaque_constant_type(type(value)):
             return value.__fx_repr__()[0]
         # repr(float('inf')) is 'inf', which is not an evaluable literal.
         if isinstance(value, float) and not math.isfinite(value):
@@ -285,7 +285,7 @@ def _collect(value: Any) -> None:
                 return
             if isinstance(value, OpaqueValueBundle.PRIMITIVE_TYPES):
                 return
-            if _is_opaque_value_type(type(value)):
+            if _is_opaque_constant_type(type(value)):
                 _, extra = value.__fx_repr__()
                 globals_.update(extra)
 
@@ -297,18 +297,18 @@ def _collect(value: Any) -> None:
 try:
     from torch._library.opaque_object import (
         get_opaque_type_name,
-        is_opaque_value_type as _is_opaque_value_type,
-        register_opaque_type,
+        is_opaque_constant_type as _is_opaque_constant_type,
+        register_custom_class,
     )
 
-    register_opaque_type(OpaqueValueBundle, typ="value")
+    register_custom_class(OpaqueValueBundle, typ="constant")
     _OPAQUE_VALUE_BUNDLE_TYPE_NAME: Optional[str] = get_opaque_type_name(OpaqueValueBundle)
 # Older torch without opaque_object support.
 except Exception as e:  # pylint: disable=broad-exception-caught  # pragma: no cover
     record_compile_disabled(
         f"could not register OpaqueValueBundle as an opaque type ({e}); use a newer PyTorch build"
     )
-    _is_opaque_value_type = None
+    _is_opaque_constant_type = None
     _OPAQUE_VALUE_BUNDLE_TYPE_NAME = None
 
 try:
@@ -461,8 +461,8 @@ def _is_simple_annot(annot: Any) -> bool:
         return True
     if (
         isinstance(annot, type)
-        and _is_opaque_value_type is not None
-        and _is_opaque_value_type(annot)
+        and _is_opaque_constant_type is not None
+        and _is_opaque_constant_type(annot)
     ):
         return True
     if get_origin(annot) in (tuple, list):

transformer_engine/pytorch/dynamo/quantizer_opaque.py

@@ -108,15 +108,15 @@ def register_value_opaque_quantizer(cls: type) -> None:
                 "the field in ``_rebuild_derived_state`` instead."
             )
     cls._value_field_names = tuple(fields)
-    # ``register_opaque_type`` requires ``__fx_repr__`` to already exist on the
+    # ``register_custom_class`` requires ``__fx_repr__`` to already exist on the
     # class, so attach it before registering.
     if "__fx_repr__" not in cls.__dict__:
         cls.__fx_repr__ = _quantizer_fx_repr
 
     try:
         from torch._library.opaque_object import (  # pylint: disable=import-outside-toplevel
-            register_opaque_type,
-            is_opaque_value_type,
+            register_custom_class,
+            is_opaque_constant_type,
         )
     except (ImportError, AttributeError) as e:
         # Older PyTorch without the opaque-object API: eager value semantics
@@ -127,8 +127,8 @@ def register_value_opaque_quantizer(cls: type) -> None:
         return
 
     try:
-        if not is_opaque_value_type(cls):
-            register_opaque_type(cls, typ="value")
+        if not is_opaque_constant_type(cls):
+            register_custom_class(cls, typ="constant")
     except (RuntimeError, TypeError) as e:
         # Keep TE importable: neither the opaque-type query nor the registration
         # must crash the import, e.g. on PyTorch versions with only partial /