[BugFix, Performance] Fewer imports at root (#682)

tensordict/memmap.py

@@ -22,20 +22,8 @@
 
 from tensordict.utils import implement_for
 
-from torch import distributed as dist
-
 from torch.multiprocessing.reductions import ForkingPickler
 
-try:
-    if dist.is_available():
-        from torch.distributed._tensor.api import DTensor
-    else:
-        raise ImportError
-except ImportError:
-
-    class DTensor(torch.Tensor):  # noqa: D101
-        ...
-
 
 class MemoryMappedTensor(torch.Tensor):
     """A Memory-mapped Tensor.
@@ -204,7 +192,7 @@ def from_tensor(
         out.index = None
         out.parent_shape = input.shape
         if copy_data:
-            if isinstance(input, DTensor):
+            if hasattr(input, "full_tensor"):
                 input = input.full_tensor()
             out.copy_(input)
         return out

tensordict/persistent.py

@@ -45,6 +45,9 @@
     NestedKey,
     NUMPY_TO_TORCH_DTYPE_DICT,
 )
+from torch import multiprocessing as mp
+
+_has_h5 = importlib.util.find_spec("h5py", None) is not None
 
 
 class _Visitor:
@@ -142,7 +145,9 @@ def __init__(
         self._locked_tensordicts = []
         self._lock_id = set()
         if not _has_h5:
-            raise ModuleNotFoundError("Could not load h5py.") from H5_ERR
+            raise ModuleNotFoundError("Could not load h5py.")
+        import h5py
+
         super().__init__()
         self.filename = filename
         self.mode = mode
@@ -204,6 +209,8 @@ def from_dict(cls, input_dict, filename, batch_size=None, device=None, **kwargs)
             A :class:`PersitentTensorDict` instance linked to the newly created file.
 
         """
+        import h5py
+
         file = h5py.File(filename, "w", locking=cls.LOCKING)
         _has_batch_size = True
         if batch_size is None:
@@ -251,6 +258,8 @@ def _get_array(self, key, default=NO_DEFAULT):
             raise KeyError(f"key {key} not found in PersistentTensorDict {self}")
 
     def _process_array(self, key, array):
+        import h5py
+
         if isinstance(array, (h5py.Dataset,)):
             if self.device is not None:
                 device = self.device
@@ -285,6 +294,8 @@ def get(self, key, default=NO_DEFAULT):
     def get_at(
         self, key: NestedKey, idx: IndexType, default: CompatibleType = NO_DEFAULT
     ) -> CompatibleType:
+        import h5py
+
         array = self._get_array(key, default)
         if isinstance(array, (h5py.Dataset,)):
             if self.device is not None:
@@ -330,6 +341,8 @@ def _get_metadata(self, key):
 
         This method avoids creating a tensor from scratch, and just reads the metadata of the array.
         """
+        import h5py
+
         array = self._get_array(key)
         if (
             isinstance(array, (h5py.Dataset,))
@@ -971,6 +984,8 @@ def _set_metadata(self, orig_metadata_container: PersistentTensorDict):
             self._nested_tensordicts[key]._set_metadata(td)
 
     def _clone(self, recurse: bool = True, newfile=None) -> PersistentTensorDict:
+        import h5py
+
         if recurse:
             # this should clone the h5 to a new location indicated by newfile
             if newfile is None:
@@ -1029,6 +1044,8 @@ def __getstate__(self):
         return state
 
     def __setstate__(self, state):
+        import h5py
+
         state["file"] = h5py.File(
             state["filename"], mode=state["mode"], locking=self.LOCKING
         )