[BugFix] Gracefully handle C++ import error in TorchRL (#1640)

Author: Vincent Moens

torchrl/_extension.py

@@ -23,3 +23,10 @@ def _init_extension():
     if not is_module_available("torchrl._torchrl"):
         warnings.warn("torchrl C++ extension is not available.")
         return
+
+
+EXTENSION_WARNING = (
+    "Failed to import torchrl C++ binaries. Some modules (eg, prioritized replay buffers) may not work with your installation. "
+    "If you installed TorchRL from PyPI, please report the bug on TorchRL github. "
+    "If you installed TorchRL locally and/or in development mode, check that you have all the required compiling packages."
+)

torchrl/data/replay_buffers/samplers.py

@@ -2,20 +2,25 @@
 #
 # This source code is licensed under the MIT license found in the
 # LICENSE file in the root directory of this source tree.
-
+import warnings
 from abc import ABC, abstractmethod
 from copy import deepcopy
 from typing import Any, Dict, Tuple, Union
 
 import numpy as np
 import torch
 
-from torchrl._torchrl import (
-    MinSegmentTreeFp32,
-    MinSegmentTreeFp64,
-    SumSegmentTreeFp32,
-    SumSegmentTreeFp64,
-)
+from ..._extension import EXTENSION_WARNING
+
+try:
+    from torchrl._torchrl import (
+        MinSegmentTreeFp32,
+        MinSegmentTreeFp64,
+        SumSegmentTreeFp32,
+        SumSegmentTreeFp64,
+    )
+except ImportError:
+    warnings.warn(EXTENSION_WARNING)
 
 from .storages import Storage
 from .utils import _to_numpy, INT_CLASSES

torchrl/modules/distributions/continuous.py

@@ -15,7 +15,6 @@
     TruncatedNormal as _TruncatedNormal,
 )
 
-# from torchrl._torchrl import safeatanh, safetanh
 from torchrl.modules.distributions.utils import (
     _cast_device,
     FasterTransformedDistribution,