[Feature] _foreach_copy_ for update_

ghstack-source-id: dc32497757d1fb19dc61dea9115810890d5a4acb Pull Request resolved: #1032
pytorch · Oct 8, 2024 · 177a08a · 177a08a
1 parent a4b04cf
commit 177a08a
Showing 1 changed file with 20 additions and 22 deletions.
diff --git a/tensordict/base.py b/tensordict/base.py
@@ -5193,30 +5193,29 @@ def inplace_update(name, dest, source):
                     if key == name[: len(key)]:
                         dest.copy_(source, non_blocking=non_blocking)
 
+            self._apply_nest(
+                inplace_update,
+                input_dict_or_td,
+                nested_keys=True,
+                default=None,
+                filter_empty=True,
+                named=named,
+                is_leaf=_is_leaf_nontensor,
+            )
+            return self
         else:
-            named = False
-
-            def inplace_update(dest, source):
-                if source is None:
-                    return None
-                dest.copy_(source, non_blocking=non_blocking)
+            if not _is_tensor_collection(type(input_dict_or_td)):
+                from tensordict import TensorDict
 
-        if not _is_tensor_collection(type(input_dict_or_td)):
-            from tensordict import TensorDict
+                input_dict_or_td = TensorDict.from_dict(
+                    input_dict_or_td, batch_dims=self.batch_dims
+                )
 
-            input_dict_or_td = TensorDict.from_dict(
-                input_dict_or_td, batch_dims=self.batch_dims
-            )
-        self._apply_nest(
-            inplace_update,
-            input_dict_or_td,
-            nested_keys=True,
-            default=None,
-            filter_empty=True,
-            named=named,
-            is_leaf=_is_leaf_nontensor,
-        )
-        return self
+            # Fastest route using _foreach_copy_
+            keys, vals = self._items_list(True, True)
+            other_val = input_dict_or_td._values_list(True, True, sorting_keys=keys)
+            torch._foreach_copy_(vals, other_val)
+            return self
 
     def update_at_(
         self,
@@ -8157,7 +8156,6 @@ def add_(self, other: TensorDictBase | float, *, alpha: float | None = None):
         .. note::
             In-place ``add`` does not support ``default`` keyword argument.
         """
-        torch.Tensor.add_
         if _is_tensor_collection(type(other)):
             keys, vals = self._items_list(True, True)
             other_val = other._values_list(True, True, sorting_keys=keys)