[Feature] torch.where (#507)

tensordict/nn/params.py

@@ -190,9 +190,7 @@ def _carry_over(func):
     @wraps(func)
     def new_func(self, *args, **kwargs):
         out = getattr(self._param_td, name)(*args, **kwargs)
-        print("out is", out)
         out = TensorDictParams(out, no_convert=True)
-        print("out is (2)", out)
         out.no_convert = self.no_convert
         return out
 
@@ -667,6 +665,10 @@ def exclude(self, *keys: str, inplace: bool = False) -> TensorDictBase:
     def transpose(self, dim0, dim1):
         ...
 
+    @_carry_over
+    def where(self, condition, other, *, out=None):
+        ...
+
     @_carry_over
     def permute(
         self,

tensordict/persistent.py

@@ -506,6 +506,9 @@ def is_contiguous(self):
     def masked_fill(self, mask, value):
         return self.to_tensordict().masked_fill(mask, value)
 
+    def where(self, condition, other, *, out=None):
+        return self.to_tensordict().where(condition=condition, other=other, out=out)
+
     def masked_fill_(self, mask, value):
         for key in self.keys(include_nested=True, leaves_only=True):
             array = self._get_array(key)

tensordict/tensordict.py

@@ -2593,6 +2593,19 @@ def masked_fill(self, mask: Tensor, value: float | bool) -> TensorDictBase:
         """
         raise NotImplementedError
 
+    def where(self, condition, other, *, out=None):
+        """Return a ``TensorDict`` of elements selected from either self or other, depending on condition.
+
+        Args:
+            condition (BoolTensor): When ``True`` (nonzero), yields ``self``,
+                otherwise yields ``other``.
+            other (TensorDictBase or Scalar): value (if ``other`` is a scalar)
+                or values selected at indices where condition is ``False``.
+            out (Tensor, optional): the output ``TensorDictBase`` instance.
+
+        """
+        raise NotImplementedError
+
     def masked_select(self, mask: Tensor) -> TensorDictBase:
         """Masks all tensors of the TensorDict and return a new TensorDict instance with similar keys pointing to masked values.
 
@@ -4431,6 +4444,42 @@ def to(tensor):
                 f"instance, {dest} not allowed"
             )
 
+    def where(self, condition, other, *, out=None):
+        if out is None:
+            if _is_tensor_collection(other.__class__):
+
+                def func(tensor, _other):
+                    return torch.where(
+                        expand_as_right(condition, tensor), tensor, _other
+                    )
+
+                return self.apply(func, other)
+            else:
+
+                def func(tensor):
+                    return torch.where(
+                        expand_as_right(condition, tensor), tensor, other
+                    )
+
+                return self.apply(func)
+        else:
+            if _is_tensor_collection(other.__class__):
+
+                def func(tensor, _other, _out):
+                    return torch.where(
+                        expand_as_right(condition, tensor), tensor, _other, out=_out
+                    )
+
+                return self.apply(func, other, out)
+            else:
+
+                def func(tensor, _out):
+                    return torch.where(
+                        expand_as_right(condition, tensor), tensor, other, out=_out
+                    )
+
+                return self.apply(func, out)
+
     def masked_fill_(self, mask: Tensor, value: float | int | bool) -> TensorDictBase:
         for item in self.values():
             mask_expand = expand_as_right(mask, item)
@@ -4807,9 +4856,7 @@ def _zeros_like(td: TensorDictBase, **kwargs: Any) -> TensorDictBase:
 
 @implements_for_td(torch.ones_like)
 def _ones_like(td: TensorDictBase, **kwargs: Any) -> TensorDictBase:
-    td_clone = td.clone()
-    for key in td_clone.keys():
-        td_clone.fill_(key, 1.0)
+    td_clone = td.apply(lambda x: torch.ones_like(x))
     if "device" in kwargs:
         td_clone = td_clone.to(kwargs.pop("device"))
     if len(kwargs):
@@ -5798,6 +5845,9 @@ def pin_memory(self) -> TensorDictBase:
     def detach_(self) -> TensorDictBase:
         raise RuntimeError("Detaching a sub-tensordict in-place cannot be done.")
 
+    def where(self, condition, other, *, out=None):
+        return self.to_tensordict().where(condition=condition, other=other, out=out)
+
     def masked_fill_(self, mask: Tensor, value: float | bool) -> TensorDictBase:
         for key, item in self.items():
             self.set_(key, torch.full_like(item, value))
@@ -7536,6 +7586,27 @@ def sort_keys(element):
 
     rename_key = _renamed_inplace_method(rename_key_)
 
+    def where(self, condition, other, *, out=None):
+        condition = condition.unbind(self.stack_dim)
+        if _is_tensor_collection(other.__class__) or (
+            isinstance(other, Tensor)
+            and other.shape[: self.stack_dim] == self.shape[: self.stack_dim]
+        ):
+            other = other.unbind(self.stack_dim)
+            return torch.stack(
+                [
+                    td.where(cond, _other)
+                    for td, cond, _other in zip(self.tensordicts, condition, other)
+                ],
+                self.stack_dim,
+                out=out,
+            )
+        return torch.stack(
+            [td.where(cond, other) for td, cond in zip(self.tensordicts, condition)],
+            self.stack_dim,
+            out=out,
+        )
+
     def masked_fill_(self, mask: Tensor, value: float | bool) -> TensorDictBase:
         mask_unbind = mask.unbind(dim=self.stack_dim)
         for _mask, td in zip(mask_unbind, self.tensordicts):
@@ -8027,6 +8098,9 @@ def detach_(self) -> _CustomOpTensorDict:
         self._source.detach_()
         return self
 
+    def where(self, condition, other, *, out=None):
+        return self.to_tensordict().where(condition=condition, other=other, out=out)
+
     def masked_fill_(self, mask: Tensor, value: float | bool) -> _CustomOpTensorDict:
         for key, item in self.items():
             val = self._source.get(key)
@@ -8892,3 +8966,23 @@ def _dispatch(remaining_index, stack_index, i=None):
         remaining_index = _dispatch(remaining_index, stack_index.tolist())
     out["remaining_index"] = _reduce_index(remaining_index)
     return out
+
+
+@implements_for_td(torch.where)
+def where(condition, input, other, *, out=None):
+    """Return a ``TensorDict`` of elements selected from either input or other, depending on condition.
+
+    Args:
+        condition (BoolTensor): When ``True`` (nonzero), yield ``input``, otherwise yield ``other``.
+        input (TensorDictBase or Scalar): value (if ``input`` is a scalar) or values selected at indices where condition is ``True``.
+        other (TensorDictBase or Scalar): value (if ``other`` is a scalar) or values selected at indices where condition is ``False``.
+        out (Tensor, optional): the output ``TensorDictBase`` instance.
+
+    """
+    from tensordict.persistent import PersistentTensorDict
+
+    if isinstance(out, PersistentTensorDict):
+        raise RuntimeError(
+            "Cannot use a persistent tensordict as output of torch.where."
+        )
+    return input.where(condition, other, out=out)

test/test_tensordict.py

@@ -1143,6 +1143,37 @@ def test_gather(self, td_name, device, dim):
         td_gather2 = torch.gather(td, dim=dim, index=index, out=out)
         assert (td_gather2 != 0).any()
 
+    def test_where(self, td_name, device):
+        torch.manual_seed(1)
+        td = getattr(self, td_name)(device)
+        mask = torch.zeros(td.shape, dtype=torch.bool, device=device).bernoulli_()
+        td_where = torch.where(mask, td, 0)
+        for k in td.keys(True, True):
+            assert (td_where.get(k)[~mask] == 0).all()
+        td_where = torch.where(mask, td, torch.ones_like(td))
+        for k in td.keys(True, True):
+            assert (td_where.get(k)[~mask] == 1).all()
+        td_where = td.clone()
+        # torch.where(mask, td, torch.zeros((), device=device), out=td_where)
+        # for k in td.keys(True, True):
+        #     assert (td_where.get(k)[~mask] == 0).all()
+        if td_name == "td_params":
+            with pytest.raises(
+                RuntimeError, match="don't support automatic differentiation"
+            ):
+                torch.where(mask, td, torch.ones_like(td), out=td_where)
+            return
+        if td_name == "td_h5":
+            with pytest.raises(
+                RuntimeError,
+                match="Cannot use a persistent tensordict as output of torch.where",
+            ):
+                torch.where(mask, td, torch.ones_like(td), out=td_where)
+            return
+        torch.where(mask, td, torch.ones_like(td), out=td_where)
+        for k in td.keys(True, True):
+            assert (td_where.get(k)[~mask] == 1).all()
+
     def test_masking_set(self, td_name, device):
         torch.manual_seed(1)
         td = getattr(self, td_name)(device)