[BE] Ensure abstractmethods are implemented for specs

ghstack-source-id: 7b943aa
Pull Request resolved: #2790

Author: Vincent Moens

torchrl/data/tensor_specs.py

@@ -540,7 +540,7 @@ def __repr__(self):
 
 
 @dataclass(repr=False)
-class TensorSpec:
+class TensorSpec(metaclass=abc.ABCMeta):
     """Parent class of the tensor meta-data containers.
 
     TorchRL's TensorSpec are used to present what input/output is to be expected for a specific class,
@@ -675,6 +675,11 @@ def encode(
             self.assert_is_in(val)
         return val
 
+    @abc.abstractmethod
+    def __eq__(self, other: Any) -> bool:
+        # Implement minimal version if super() is called
+        return type(self) is type(other)
+
     def __ne__(self, other):
         return not (self == other)
 
@@ -734,13 +739,31 @@ def index(
     ) -> torch.Tensor | TensorDictBase:
         """Indexes the input tensor.
 
+        This method is to be used with specs that encode one or more categorical variables (e.g.,
+        :class:`~torchrl.data.OneHot` or :class:`~torchrl.data.Categorical`), such that indexing of a tensor
+        with a sample can be done without caring about the actual representation of the index.
+
         Args:
             index (int, torch.Tensor, slice or list): index of the tensor
             tensor_to_index: tensor to be indexed
 
         Returns:
             indexed tensor
 
+        Exanples:
+            >>> from torchrl.data import OneHot
+            >>> import torch
+            >>>
+            >>> one_hot = OneHot(n=100)
+            >>> categ = one_hot.to_categorical_spec()
+            >>> idx_one_hot = torch.zeros((100,), dtype=torch.bool)
+            >>> idx_one_hot[50] = 1
+            >>> print(one_hot.index(idx_one_hot, torch.arange(100)))
+            tensor(50)
+            >>> idx_categ = one_hot.to_categorical(idx_one_hot)
+            >>> print(categ.index(idx_categ, torch.arange(100)))
+            tensor(50)
+
         """
         ...
 
@@ -1306,6 +1329,31 @@ class Stacked(_LazyStackedMixin[TensorSpec], TensorSpec):
 
     """
 
+    def _reshape(
+        self,
+        *args,
+        **kwargs,
+    ) -> Any:
+        raise NotImplementedError(
+            f"`reshape` is not implemented for {type(self).__name__} specs."
+        )
+
+    def cardinality(
+        self,
+        *args,
+        **kwargs,
+    ) -> Any:
+        raise NotImplementedError(
+            f"`cardinality` is not implemented for {type(self).__name__} specs."
+        )
+
+    def index(
+        self, index: INDEX_TYPING, tensor_to_index: torch.Tensor | TensorDictBase
+    ) -> torch.Tensor | TensorDictBase:
+        raise NotImplementedError(
+            f"`index` is not implemented for {type(self).__name__} specs."
+        )
+
     def __eq__(self, other):
         if not isinstance(other, Stacked):
             return False
@@ -1829,7 +1877,7 @@ def index(self, index: INDEX_TYPING, tensor_to_index: torch.Tensor) -> torch.Ten
                 f"Only tensors are allowed for indexing using "
                 f"{self.__class__.__name__}.index(...)"
             )
-        index = index.nonzero().squeeze()
+        index = index.nonzero(as_tuple=True)[-1]
         index = index.expand((*tensor_to_index.shape[:-1], index.shape[-1]))
         return tensor_to_index.gather(-1, index)
 
@@ -2148,6 +2196,11 @@ def __init__(
             domain=domain,
         )
 
+    def index(
+        self, index: INDEX_TYPING, tensor_to_index: torch.Tensor | TensorDictBase
+    ) -> torch.Tensor | TensorDictBase:
+        raise NotImplementedError("Indexing not implemented for Bounded.")
+
     def enumerate(self, use_mask: bool = False) -> Any:
         raise NotImplementedError(
             f"enumerate is not implemented for spec of class {type(self).__name__}."
@@ -2484,11 +2537,19 @@ def __eq__(self, other):
         eq = eq & (self.example_data == getattr(other, "example_data", None))
         return eq
 
+    def _project(self) -> Any:
+        raise NotImplementedError("Cannot project a NonTensorSpec.")
+
+    def index(
+        self, index: INDEX_TYPING, tensor_to_index: torch.Tensor | TensorDictBase
+    ) -> torch.Tensor | TensorDictBase:
+        raise NotImplementedError("Cannot use index with a NonTensorSpec.")
+
     def cardinality(self) -> Any:
-        raise RuntimeError("Cannot enumerate a NonTensorSpec.")
+        raise NotImplementedError("Cannot enumerate a NonTensorSpec.")
 
     def enumerate(self, use_mask: bool = False) -> Any:
-        raise RuntimeError("Cannot enumerate a NonTensorSpec.")
+        raise NotImplementedError("Cannot enumerate a NonTensorSpec.")
 
     def to(self, dest: Union[torch.dtype, DEVICE_TYPING]) -> NonTensor:
         if isinstance(dest, torch.dtype):
@@ -2752,6 +2813,16 @@ def __init__(
             shape=shape, space=box, device=device, dtype=dtype, domain=domain, **kwargs
         )
 
+    def cardinality(self) -> int:
+        raise NotImplementedError(
+            "`cardinality` is not implemented for Unbounded specs."
+        )
+
+    def index(
+        self, index: INDEX_TYPING, tensor_to_index: torch.Tensor | TensorDictBase
+    ) -> torch.Tensor | TensorDictBase:
+        raise NotImplementedError("`index` is not implemented for Unbounded specs.")
+
     def to(self, dest: Union[torch.dtype, DEVICE_TYPING]) -> Unbounded:
         if isinstance(dest, torch.dtype):
             dest_dtype = dest
@@ -3527,6 +3598,14 @@ def rand(self, shape: torch.Size = None) -> torch.Tensor:
         out = torch.multinomial(mask_flat.float(), 1).reshape(shape_out)
         return out
 
+    def index(
+        self, index: INDEX_TYPING, tensor_to_index: torch.Tensor | TensorDictBase
+    ) -> torch.Tensor | TensorDictBase:
+        idx = index.expand(
+            tensor_to_index.shape[: -self.ndim] + torch.Size([-1] * self.ndim)
+        )
+        return tensor_to_index.gather(-1, idx)
+
     def _project(self, val: torch.Tensor) -> torch.Tensor:
         if val.dtype not in (torch.int, torch.long):
             val = torch.round(val)
@@ -3863,9 +3942,50 @@ def cardinality(self) -> int:
                 .item()
             )
 
+    def enumerate(self, use_mask: bool = False) -> List[Any]:
+        return [s for choice in self._choices for s in choice.enumerate()]
+
+    def _project(
+        self, val: torch.Tensor | TensorDictBase
+    ) -> torch.Tensor | TensorDictBase:
+        raise NotImplementedError(
+            "_project is not implemented for Choice. If this feature is required, please raise "
+            "an issue on TorchRL github repo."
+        )
+
+    def _reshape(self, shape: torch.Size) -> T:
+        return self.__class__(
+            [choice.reshape(shape) for choice in self._choices],
+        )
+
+    def index(
+        self, index: INDEX_TYPING, tensor_to_index: torch.Tensor | TensorDictBase
+    ) -> torch.Tensor | TensorDictBase:
+        raise NotImplementedError(
+            "index is not implemented for Choice. If this feature is required, please raise "
+            "an issue on TorchRL github repo."
+        )
+
+    @property
+    def num_choices(self):
+        """Number of choices for the spec."""
+        return len(self._choices)
+
     def to(self, dest: Union[torch.dtype, DEVICE_TYPING]) -> Choice:
         return self.__class__([choice.to(dest) for choice in self._choices])
 
+    def __eq__(self, other):
+        if not isinstance(other, Choice):
+            return False
+        if self.num_choices != other.num_choices:
+            return False
+        return all(
+            (s0 == s1).all()
+            if isinstance(s0, torch.Tensor) or is_tensor_collection(s0)
+            else s0 == s1
+            for s0, s1 in zip(self._choices, other._choices)
+        )
+
 
 @dataclass(repr=False)
 class Binary(Categorical):
@@ -4643,6 +4763,24 @@ def shape(self, value: torch.Size):
                     )
         self._shape = _size(value)
 
+    def _project(
+        self, val: torch.Tensor | TensorDictBase
+    ) -> torch.Tensor | TensorDictBase:
+        if self.data_cls is None:
+            cls = TensorDict
+        else:
+            cls = self.data_cls
+        return cls.from_dict(
+            {k: item._project(val[k]) for k, item in self.items()},
+            batch_size=self.shape,
+            device=self.device,
+        )
+
+    def index(
+        self, index: INDEX_TYPING, tensor_to_index: torch.Tensor | TensorDictBase
+    ) -> torch.Tensor | TensorDictBase:
+        raise NotImplementedError("`index` is not implemented for Composite specs.")
+
     def is_empty(self, recurse: bool = False):
         """Whether the composite spec contains specs or not.
 
@@ -5569,6 +5707,31 @@ class StackedComposite(_LazyStackedMixin[Composite], Composite):
 
     """
 
+    def _reshape(
+        self,
+        *args,
+        **kwargs,
+    ) -> Any:
+        raise NotImplementedError(
+            f"`reshape` is not implemented for {type(self).__name__} specs."
+        )
+
+    def cardinality(
+        self,
+        *args,
+        **kwargs,
+    ) -> Any:
+        raise NotImplementedError(
+            f"`cardinality` is not implemented for {type(self).__name__} specs."
+        )
+
+    def index(
+        self, index: INDEX_TYPING, tensor_to_index: torch.Tensor | TensorDictBase
+    ) -> torch.Tensor | TensorDictBase:
+        raise NotImplementedError(
+            f"`index` is not implemented for {type(self).__name__} specs."
+        )
+
     def update(self, dict) -> None:
         for key, item in dict.items():
             if key in self.keys() and isinstance(