diff --git a/tests/pytorch/test_torch_compile.py b/tests/pytorch/test_torch_compile.py index f8e09d5ce1..e1c6eaf599 100644 --- a/tests/pytorch/test_torch_compile.py +++ b/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, diff --git a/transformer_engine/pytorch/dynamo/custom_op.py b/transformer_engine/pytorch/dynamo/custom_op.py index 00846d615a..f887a48d95 100644 --- a/transformer_engine/pytorch/dynamo/custom_op.py +++ b/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): diff --git a/transformer_engine/pytorch/dynamo/quantizer_opaque.py b/transformer_engine/pytorch/dynamo/quantizer_opaque.py index dc689258d1..da595b1070 100644 --- a/transformer_engine/pytorch/dynamo/quantizer_opaque.py +++ b/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 /