[Performance] Faster update_ (#705)

tensordict/_lazy.py

@@ -1413,6 +1413,7 @@ def _apply_nest(
         nested_keys: bool = False,
         prefix: tuple = (),
         filter_empty: bool | None = None,
+        is_leaf: Callable | None = None,
         **constructor_kwargs,
     ) -> T | None:
         if inplace and any(
@@ -1438,6 +1439,7 @@ def _apply_nest(
                 prefix=prefix,
                 inplace=inplace,
                 filter_empty=filter_empty,
+                is_leaf=is_leaf,
                 **constructor_kwargs,
             )
 
@@ -1455,6 +1457,7 @@ def _apply_nest(
                 prefix=prefix,  # + (i,),
                 inplace=inplace,
                 filter_empty=filter_empty,
+                is_leaf=is_leaf,
             )
             for i, (td, *oth) in enumerate(zip(self.tensordicts, *others))
         ]

tensordict/_td.py

@@ -661,6 +661,7 @@ def _apply_nest(
         nested_keys: bool = False,
         prefix: tuple = (),
         filter_empty: bool | None = None,
+        is_leaf: Callable = None,
         **constructor_kwargs,
     ) -> T | None:
         if inplace:
@@ -696,10 +697,13 @@ def make_result():
             is_locked = False
 
         any_set = False
+        if is_leaf is None:
+            is_leaf = _default_is_leaf
+
         for key, item in self.items():
             if (
                 not call_on_nested
-                and _is_tensor_collection(item.__class__)
+                and not is_leaf(item.__class__)
                 # and not is_non_tensor(item)
             ):
                 if default is not NO_DEFAULT:
@@ -725,6 +729,7 @@ def make_result():
                     default=default,
                     prefix=prefix + (key,),
                     filter_empty=filter_empty,
+                    is_leaf=is_leaf,
                     **constructor_kwargs,
                 )
             else: