[Refactor] Version-gate HigherOrderOperator dispatch with implement_for

tensordict/_torch_func.py

@@ -30,6 +30,7 @@
     _shape,
     _zip_strict,
     DeviceType,
+    implement_for,
     is_tensorclass,
     lazy_legacy,
     set_lazy_legacy,
@@ -60,6 +61,31 @@ def tree_leaves(pytree):
         return tree_flatten(pytree)[0]
 
 
+try:
+    from torch._ops import HigherOrderOperator as _HigherOrderOperator
+except ImportError:
+    _HigherOrderOperator = None  # torch < 2.1
+
+
+@implement_for("torch", "2.1")
+def _maybe_dispatch_higher_order_op(func, args, kwargs):
+    """Pass through HigherOrderOperator calls (e.g. invoke_subgraph) to avoid returning NotImplemented."""
+    if isinstance(func, _HigherOrderOperator):
+        with torch._C.DisableTorchFunctionSubclass():
+            return func(*args, **kwargs)
+    return NotImplemented
+
+
+@implement_for("torch", None, "2.1")
+def _maybe_dispatch_higher_order_op(func, args, kwargs):  # noqa: F811
+    return NotImplemented
+
+
+# Force implement_for resolution at import time so that it doesn't trigger
+# _gcd_import inside torch.compile tracing context.
+_maybe_dispatch_higher_order_op(None, (), {})
+
+
 def implements_for_td(torch_function: Callable) -> Callable[[Callable], Callable]:
     """Register a torch function override for TensorDict."""
 

tensordict/base.py

@@ -711,19 +711,17 @@ def __torch_function__(
         args: tuple[Any, ...] = (),
         kwargs: dict[str, Any] | None = None,
     ) -> Callable:
-        from tensordict._torch_func import TD_HANDLED_FUNCTIONS
+        from tensordict._torch_func import (
+            _maybe_dispatch_higher_order_op,
+            TD_HANDLED_FUNCTIONS,
+        )
 
         if kwargs is None:
             kwargs = {}
         if func not in TD_HANDLED_FUNCTIONS or not all(
             issubclass(t, (Tensor, TensorDictBase)) or _is_tensorclass(t) for t in types
         ):
-            from torch._ops import HigherOrderOperator
-
-            if isinstance(func, HigherOrderOperator):
-                with torch._C.DisableTorchFunctionSubclass():
-                    return func(*args, **kwargs)
-            return NotImplemented
+            return _maybe_dispatch_higher_order_op(func, args, kwargs)
         return TD_HANDLED_FUNCTIONS[func](*args, **kwargs)
 
     @abc.abstractmethod

tensordict/nn/params.py

@@ -31,7 +31,7 @@
 from tensordict._nestedkey import NestedKey
 from tensordict._td import _SubTensorDict, TensorDict
 from tensordict._tensorcollection import TensorCollection
-from tensordict._torch_func import TD_HANDLED_FUNCTIONS
+from tensordict._torch_func import _maybe_dispatch_higher_order_op, TD_HANDLED_FUNCTIONS
 
 from tensordict.base import (
     _default_is_leaf,
@@ -512,12 +512,7 @@ def __torch_function__(
         if func not in TDPARAM_HANDLED_FUNCTIONS or not all(
             issubclass(t, (Tensor, ftdim.Tensor, TensorDictBase)) for t in types
         ):
-            from torch._ops import HigherOrderOperator
-
-            if isinstance(func, HigherOrderOperator):
-                with torch._C.DisableTorchFunctionSubclass():
-                    return func(*args, **kwargs)
-            return NotImplemented
+            return _maybe_dispatch_higher_order_op(func, args, kwargs)
         return TDPARAM_HANDLED_FUNCTIONS[func](*args, **kwargs)
 
     @classmethod

tensordict/tensorclass.py

@@ -46,7 +46,7 @@
 from tensordict._pytree import _register_td_node
 from tensordict._td import is_tensor_collection, NO_DEFAULT, TensorDict, TensorDictBase
 from tensordict._tensorcollection import TensorCollection
-from tensordict._torch_func import TD_HANDLED_FUNCTIONS
+from tensordict._torch_func import _maybe_dispatch_higher_order_op, TD_HANDLED_FUNCTIONS
 from tensordict.base import (
     _ACCEPTED_CLASSES,
     _GET_DEFAULTS_TO_NONE,
@@ -943,14 +943,9 @@ def __torch_function__(
         if func not in _TD_PASS_THROUGH or not all(
             issubclass(t, (Tensor, cls, TensorDictBase)) for t in types
         ):
-            from torch._ops import HigherOrderOperator
-
-            if isinstance(func, HigherOrderOperator):
-                if kwargs is None:
-                    kwargs = {}
-                with torch._C.DisableTorchFunctionSubclass():
-                    return func(*args, **kwargs)
-            return NotImplemented
+            if kwargs is None:
+                kwargs = {}
+            return _maybe_dispatch_higher_order_op(func, args, kwargs)
 
         if kwargs is None:
             kwargs = {}
@@ -3906,14 +3901,9 @@ def __torch_function__(
         if func not in _TD_PASS_THROUGH or not all(
             issubclass(t, (Tensor, cls)) for t in types
         ):
-            from torch._ops import HigherOrderOperator
-
-            if isinstance(func, HigherOrderOperator):
-                if kwargs is None:
-                    kwargs = {}
-                with torch._C.DisableTorchFunctionSubclass():
-                    return func(*args, **kwargs)
-            return NotImplemented
+            if kwargs is None:
+                kwargs = {}
+            return _maybe_dispatch_higher_order_op(func, args, kwargs)
 
         escape_conversion = func in (torch.stack,)