[BugFix] Better prop of args in update

ghstack-source-id: 38859ac
Pull Request resolved: #1290

Author: Vincent Moens

tensordict/_lazy.py

@@ -30,20 +30,7 @@
 
 import numpy as np
 
-import orjson as json
 import torch
-
-from tensordict.memmap import MemoryMappedTensor
-from torch.nn.utils.rnn import pad_sequence
-
-try:
-    from functorch import dim as ftdim
-
-    _has_funcdim = True
-except ImportError:
-    from tensordict.utils import _ftdim_mock as ftdim
-
-    _has_funcdim = False
 from tensordict._td import _SubTensorDict, _TensorDictKeysView, TensorDict
 from tensordict.base import (
     _is_leaf_nontensor,
@@ -58,6 +45,8 @@
     T,
     TensorDictBase,
 )
+
+from tensordict.memmap import MemoryMappedTensor
 from tensordict.utils import (
     _as_context_manager,
     _broadcast_tensors,
@@ -90,7 +79,22 @@
     unravel_key_list,
 )
 from torch import Tensor
+from torch.nn.utils.rnn import pad_sequence
 
+try:
+    import orjson as json
+except ImportError:
+    # Fallback
+    import json
+
+try:
+    from functorch import dim as ftdim
+
+    _has_funcdim = True
+except ImportError:
+    from tensordict.utils import _ftdim_mock as ftdim
+
+    _has_funcdim = False
 
 _has_functorch = False
 try:
@@ -2943,6 +2947,7 @@ def update(
                     keys_to_update=keys_to_update,
                     non_blocking=non_blocking,
                     is_leaf=is_leaf,
+                    update_batch_size=update_batch_size,
                     **kwargs,
                 )
             return self

tensordict/_td.py

@@ -18,7 +18,7 @@
 from warnings import warn
 
 import numpy as np
-import orjson as json
+
 import torch
 from tensordict._nestedkey import NestedKey
 
@@ -91,6 +91,11 @@
 from torch.nn.utils._named_member_accessor import swap_tensor
 from torch.utils._pytree import tree_map
 
+try:
+    import orjson as json
+except ImportError:
+    # Fallback
+    import json
 try:
     from functorch import dim as ftdim
 

tensordict/base.py

@@ -46,7 +46,7 @@
 )
 
 import numpy as np
-import orjson as json
+
 import torch
 
 from tensordict._contextlib import LAST_OP_MAPS
@@ -110,6 +110,12 @@
 from torch.nn.parameter import Parameter, UninitializedTensorMixin
 from torch.utils._pytree import tree_map
 
+try:
+    import orjson as json
+except ImportError:
+    # Fallback for 3.13
+    import json
+
 try:
     from torch.compiler import is_compiling
 except ImportError:  # torch 2.0

tensordict/persistent.py

@@ -18,7 +18,6 @@
 
 import numpy as np
 
-import orjson as json
 import torch
 
 from tensordict._td import (
@@ -57,6 +56,12 @@
 )
 from torch import multiprocessing as mp
 
+try:
+    import orjson as json
+except ImportError:
+    # Fallback for 3.13
+    import json
+
 _has_h5 = importlib.util.find_spec("h5py", None) is not None
 
 

tensordict/tensorclass.py

@@ -41,7 +41,7 @@
 from warnings import warn
 
 import numpy as np
-import orjson as json
+
 import tensordict as tensordict_lib
 
 import torch
@@ -53,6 +53,7 @@
 from tensordict.base import (
     _ACCEPTED_CLASSES,
     _GET_DEFAULTS_TO_NONE,
+    _is_leaf_nontensor,
     _is_tensor_collection,
     _register_tensor_class,
     CompatibleType,
@@ -79,6 +80,11 @@
 from torch.multiprocessing import Manager
 from torch.utils._pytree import tree_map
 
+try:
+    import orjson as json
+except ImportError:
+    # Fallback for 3.13
+    import json
 try:
     from torch.compiler import is_compiling
 except ImportError:  # torch 2.0
@@ -1841,7 +1847,10 @@ def _update(
     non_blocking: bool = False,
     update_batch_size: bool = False,
     ignore_lock: bool = False,
+    is_leaf: Callable[[Type], bool] | None = None,
 ):
+    if is_leaf is None:
+        is_leaf = _is_leaf_nontensor
     if isinstance(input_dict_or_td, dict):
         input_dict_or_td = self.from_dict(input_dict_or_td, auto_batch_size=False)
 
@@ -1859,6 +1868,7 @@ def _update(
             non_blocking=non_blocking,
             update_batch_size=update_batch_size,
             ignore_lock=ignore_lock,
+            is_leaf=is_leaf,
         )
         self._non_tensordict.update(non_tensordict)
         return self
@@ -1871,6 +1881,7 @@ def _update(
         non_blocking=non_blocking,
         update_batch_size=update_batch_size,
         ignore_lock=ignore_lock,
+        is_leaf=is_leaf,
     )
     # We also need to remove things from non_tensordict
     if self._non_tensordict:

tensordict/tensorclass.pyi

@@ -1119,3 +1119,6 @@ class TensorClass:
     def uint32(self) -> T: ...
     def uint64(self) -> T: ...
     def uint8(self) -> T: ...
+
+class NonTensorData(TensorClass): ...
+class NonTensorStack(TensorDictBase): ...