tensor_specs.py

# Copyright (c) Meta Platforms, Inc. and affiliates.
#
# This source code is licensed under the MIT license found in the
# LICENSE file in the root directory of this source tree.

from __future__ import annotations

import abc
import enum
import functools
import gc
import math
import warnings
import weakref
from collections.abc import Iterable
from copy import deepcopy
from dataclasses import dataclass, field
from functools import wraps
from textwrap import indent
from typing import (
    Any,
    Callable,
    Dict,
    Generic,
    List,
    overload,
    Sequence,
    Tuple,
    TypeVar,
    Union,
)

import numpy as np

import tensordict
import torch
from tensordict import (
    is_tensor_collection,
    LazyStackedTensorDict,
    NonTensorData,
    NonTensorStack,
    set_capture_non_tensor_stack,
    TensorDict,
    TensorDictBase,
    unravel_key,
)
from tensordict.base import NO_DEFAULT
from tensordict.utils import (
    _getitem_batch_size,
    expand_as_right,
    is_non_tensor,
    NestedKey,
)
from torchrl._utils import _make_ordinal_device, get_binary_env_var, implement_for

try:
    from torch.compiler import is_compiling
except ImportError:
    from torch._dynamo import is_compiling

DEVICE_TYPING = Union[torch.device, str, int]

INDEX_TYPING = Union[int, torch.Tensor, np.ndarray, slice, List]

SHAPE_INDEX_TYPING = Union[
    int,
    range,
    List[int],
    np.ndarray,
    slice,
    None,
    torch.Tensor,
    type(...),
    Tuple[
        int,
        range,
        List[int],
        np.ndarray,
        slice,
        None,
        torch.Tensor,
        type(...),
        Tuple[Any],
    ],
]

# By default, we do not check that an obs is in the domain. THis should be done when validating the env beforehand
_CHECK_SPEC_ENCODE = get_binary_env_var("CHECK_SPEC_ENCODE")

_DEFAULT_SHAPE = torch.Size((1,))

DEVICE_ERR_MSG = "device of empty Composite is not defined."
NOT_IMPLEMENTED_ERROR = NotImplementedError(
    "method is not currently implemented."
    " If you are interested in this feature please submit"
    " an issue at https://github.com/pytorch/rl/issues"
)


def _size(list_of_ints):
    # ensures that np int64 elements don't slip through Size
    # see https://github.com/pytorch/pytorch/issues/127194
    return torch.Size([int(i) for i in list_of_ints])


# Akin to TD's NO_DEFAULT but won't raise a KeyError when found in a TD or used as default
class _NoDefault(enum.IntEnum):
    ZERO = 0
    ONE = 1


NO_DEFAULT_RL = _NoDefault.ONE


def _default_dtype_and_device(
    dtype: None | torch.dtype,
    device: None | str | int | torch.device,
    allow_none_device: bool = False,
) -> tuple[torch.dtype, torch.device | None]:
    if dtype is None:
        dtype = torch.get_default_dtype()
    if device is not None:
        device = _make_ordinal_device(torch.device(device))
    elif not allow_none_device:
        device = torch.zeros(()).device
    return dtype, device


def _validate_idx(shape: list[int], idx: int, axis: int = 0):
    """Raise an IndexError if idx is out of bounds for shape[axis].

    Args:
        shape (list[int]): Input shape
        idx (int): Index, may be negative
        axis (int): Shape axis to check
    """
    if shape[axis] >= 0 and (idx >= shape[axis] or idx < 0 and -idx > shape[axis]):
        raise IndexError(
            f"index {idx} is out of bounds for axis {axis} with size {shape[axis]}"
        )


def _validate_iterable(
    idx: Iterable[Any], expected_type: type, iterable_classname: str
):
    """Raise an IndexError if the iterable contains a type different from the expected type or Iterable.

    Args:
        idx (Iterable[Any]): Iterable, may contain nested iterables
        expected_type (Type): Required item type in the Iterable (e.g. int)
        iterable_classname (str): Iterable type as a string (e.g. 'List'). Logging purpose only.
    """
    for item in idx:
        if isinstance(item, Iterable):
            _validate_iterable(item, expected_type, iterable_classname)
        else:
            if not isinstance(item, expected_type):
                raise IndexError(
                    f"{iterable_classname} indexing expects {expected_type} indices"
                )


def _slice_indexing(shape: list[int], idx: slice) -> list[int]:
    """Given an input shape and a slice index, returns the new indexed shape.

    Args:
        shape (list[int]): Input shape
        idx (slice): Index
    Returns:
        Indexed shape
    Examples:
        >>> _slice_indexing([3, 4], slice(None, 2))
        [2, 4]
        >>> list(torch.rand(3, 4)[:2].shape)
        [2, 4]
    """
    if idx.step == 0:
        raise ValueError("slice step cannot be zero")
    # Slicing an empty shape returns the shape
    if len(shape) == 0:
        return shape

    if idx.start is None:
        start = 0
    else:
        start = idx.start if idx.start >= 0 else max(shape[0] + idx.start, 0)

    if idx.stop is None:
        stop = shape[0]
    else:
        stop = idx.stop if idx.stop >= 0 else max(shape[0] + idx.stop, 0)

    step = 1 if idx.step is None else idx.step
    if step > 0:
        if start >= stop:
            n_items = 0
        else:
            stop = min(stop, shape[0])
            n_items = math.ceil((stop - start) / step)
    else:
        if start <= stop:
            n_items = 0
        else:
            start = min(start, shape[0] - 1)
            n_items = math.ceil((stop - start) / step)
    return [n_items] + shape[1:]


def _shape_indexing(
    shape: list[int] | torch.Size | tuple[int], idx: SHAPE_INDEX_TYPING
) -> list[int]:
    """Given an input shape and an index, returns the size of the resulting indexed spec.

    This function includes indexing checks and may raise IndexErrors.

    Args:
        shape (list[int], torch.Size, Tuple[int): Input shape
        idx (SHAPE_INDEX_TYPING): Index
    Returns:
        Shape of the resulting spec
    Examples:
        >>> idx = (2, ..., None)
        >>> Categorical(2, shape=(3, 4))[idx].shape
        torch.Size([4, 1])
        >>> _shape_indexing([3, 4], idx)
        torch.Size([4, 1])
    """
    if not isinstance(shape, list):
        shape = list(shape)

    if idx is Ellipsis or (
        isinstance(idx, slice) and (idx.step is idx.start is idx.stop is None)
    ):
        return shape

    if idx is None:
        return [1] + shape

    if len(shape) == 0 and (
        isinstance(idx, int)
        or isinstance(idx, range)
        or isinstance(idx, list)
        and len(idx) > 0
    ):
        raise IndexError(
            f"cannot use integer indices on 0-dimensional shape. `{idx}` received"
        )

    if isinstance(idx, int):
        _validate_idx(shape, idx)
        return shape[1:]

    if isinstance(idx, range):
        if len(idx) > 0 and (idx.start >= shape[0] or idx.stop > shape[0]):
            raise IndexError(f"index out of bounds for axis 0 with size {shape[0]}")
        return [len(idx)] + shape[1:]

    if isinstance(idx, slice):
        return _slice_indexing(shape, idx)

    if isinstance(idx, tuple):
        # Supports int, None, slice and ellipsis indices
        # Index on the current shape dimension
        shape_idx = 0
        none_dims = 0
        ellipsis = False
        prev_is_list = False
        shape_len = len(shape)
        for item_idx, item in enumerate(idx):
            if item is None:
                shape = shape[:shape_idx] + [1] + shape[shape_idx:]
                shape_idx += 1
                none_dims += 1
            elif isinstance(item, int):
                _validate_idx(shape, item, shape_idx)
                del shape[shape_idx]
            elif isinstance(item, slice):
                shape[shape_idx] = _slice_indexing([shape[shape_idx]], item)[0]
                shape_idx += 1
            elif item is Ellipsis:
                if ellipsis:
                    raise IndexError("an index can only have a single ellipsis (`...`)")
                # Move to the end of the shape, subtracted by the number of future indices impacting the dimensions (i.e. all except None and ...)
                shape_idx = len(shape) - len(
                    [i for i in idx[item_idx + 1 :] if not (i is None or i is Ellipsis)]
                )
                ellipsis = True
            elif any(
                isinstance(item, _type)
                for _type in [list, tuple, range, np.ndarray, torch.Tensor]
            ):
                while isinstance(idx, tuple) and len(idx) == 1:
                    idx = idx[0]

                # Nested tuples are handled as a list. Numpy behavior
                if isinstance(item, tuple):
                    item = list(item)

                if prev_is_list and isinstance(item, list):
                    del shape[shape_idx]
                    continue

                if isinstance(item, list):
                    prev_is_list = True

                if shape_idx >= len(shape):
                    raise IndexError("Raise IndexError: too many indices for array")

                res = _shape_indexing([shape[shape_idx]], item)
                shape = shape[:shape_idx] + res + shape[shape_idx + 1 :]
                shape_idx += len(res)
            else:
                raise IndexError(
                    f"tuple indexing only supports integers, ranges, slices (`:`), ellipsis (`...`), new axis (`None`), tuples, list, tensor and ndarray indices. {str(type(idx))} received"
                )

        if len(idx) - none_dims - int(ellipsis) > shape_len:
            raise IndexError(
                f"shape is {shape_len}-dimensional, but {len(idx) - none_dims - int(ellipsis)} dimensions were indexed"
            )
        return shape

    if isinstance(idx, list):
        # int indexing only
        _validate_iterable(idx, int, "list")
        for item in np.array(idx).reshape(-1):
            _validate_idx(shape, item, 0)
        return list(np.array(idx).shape) + shape[1:]

    if isinstance(idx, np.ndarray) or isinstance(idx, torch.Tensor):
        # Out of bounds check
        for item in idx.reshape(-1):
            _validate_idx(shape, item)
        return list(_getitem_batch_size(shape, idx))


class invertible_dict(dict):
    """An invertible dictionary.

    Examples:
        >>> my_dict = invertible_dict(a=3, b=2)
        >>> inv_dict = my_dict.invert()
        >>> assert {2, 3} == set(inv_dict.keys())
    """

    def __init__(self, *args, inv_dict=None, **kwargs):
        if inv_dict is None:
            inv_dict = {}
        super().__init__(*args, **kwargs)
        self.inv_dict = inv_dict

    def __setitem__(self, k, v):
        if v in self.inv_dict or k in self:
            raise Exception("overwriting in invertible_dict is not permitted")
        self.inv_dict[v] = k
        return super().__setitem__(k, v)

    def update(self, d):
        raise NotImplementedError

    def invert(self):
        d = invertible_dict()
        for k, value in self.items():
            d[value] = k
        return d

    def inverse(self):
        return self.inv_dict


class Box:
    """A box of values."""

    def __iter__(self):
        raise NotImplementedError

    def to(self, dest: torch.dtype | DEVICE_TYPING) -> ContinuousBox:
        raise NotImplementedError

    def __repr__(self):
        return f"{self.__class__.__name__}()"

    def clone(self) -> CategoricalBox:
        return deepcopy(self)


@dataclass(repr=False)
class ContinuousBox(Box):
    """A continuous box of values, in between a minimum (self.low) and a maximum (self.high)."""

    _low: torch.Tensor
    _high: torch.Tensor
    device: torch.device | None = None

    # We store the tensors on CPU to avoid overloading CUDA with tensors that are rarely used.
    @property
    def low(self):
        low = self._low
        if self.device is not None and low.device != self.device:
            low = low.to(self.device)
        return low

    @property
    def high(self):
        high = self._high
        if self.device is not None and high.device != self.device:
            high = high.to(self.device)
        return high

    def unbind(self, dim: int = 0):
        return tuple(
            type(self)(low, high, self.device)
            for (low, high) in zip(self.low.unbind(dim), self.high.unbind(dim))
        )

    @low.setter
    def low(self, value):
        self.device = value.device
        self._low = value

    @high.setter
    def high(self, value):
        self.device = value.device
        self._high = value

    def __post_init__(self):
        self.low = self.low.clone()
        self.high = self.high.clone()

    def __iter__(self):
        yield self.low
        yield self.high

    def to(self, dest: torch.dtype | DEVICE_TYPING) -> ContinuousBox:
        return self.__class__(self.low.to(dest), self.high.to(dest))

    def clone(self) -> ContinuousBox:
        return self.__class__(self.low.clone(), self.high.clone())

    def __repr__(self):
        min_str = indent(
            f"\nlow=Tensor(shape={self.low.shape}, device={self.low.device}, dtype={self.low.dtype}, contiguous={self.high.is_contiguous()})",
            " " * 4,
        )
        max_str = indent(
            f"\nhigh=Tensor(shape={self.high.shape}, device={self.high.device}, dtype={self.high.dtype}, contiguous={self.high.is_contiguous()})",
            " " * 4,
        )
        return f"{self.__class__.__name__}({min_str},{max_str})"

    def __eq__(self, other):
        if other is None:

            minval, maxval = _minmax_dtype(self.low.dtype)
            minval = torch.as_tensor(minval).to(self.low.device, self.low.dtype)
            maxval = torch.as_tensor(maxval).to(self.low.device, self.low.dtype)
            if (
                torch.isclose(self.low, minval).all()
                and torch.isclose(self.high, maxval).all()
            ):
                return True
            if (
                not torch.isfinite(self.low).any()
                and not torch.isfinite(self.high).any()
            ):
                return True
            return False
        return (
            type(self) == type(other)
            and self.low.dtype == other.low.dtype
            and self.high.dtype == other.high.dtype
            and self.device == other.device
            and torch.isclose(self.low, other.low).all()
            and torch.isclose(self.high, other.high).all()
        )


@dataclass(repr=False, frozen=True)
class CategoricalBox(Box):
    """A box of discrete, categorical values."""

    n: int
    register = invertible_dict()

    def __post_init__(self):
        # n could be a numpy array or a tensor, making compile go a bit crazy
        # We want to make sure we're working with a regular integer
        self.__dict__["n"] = int(self.n)

    def to(self, dest: torch.dtype | DEVICE_TYPING) -> CategoricalBox:
        return deepcopy(self)

    def __repr__(self):
        return f"{self.__class__.__name__}(n={self.n})"


class DiscreteBox(CategoricalBox):
    """Deprecated version of :class:`CategoricalBox`."""

    ...


@dataclass(repr=False)
class BoxList(Box):
    """A box of discrete values."""

    boxes: list

    def to(self, dest: torch.dtype | DEVICE_TYPING) -> BoxList:
        return BoxList([box.to(dest) for box in self.boxes])

    def __iter__(self):
        yield from self.boxes

    def __repr__(self):
        return f"{self.__class__.__name__}(boxes={self.boxes})"

    def __len__(self):
        return len(self.boxes)

    @staticmethod
    def from_nvec(nvec: torch.Tensor):
        if nvec.ndim == 0:
            return CategoricalBox(nvec.item())
        else:
            return BoxList([BoxList.from_nvec(n) for n in nvec.unbind(-1)])


@dataclass(repr=False, frozen=True)
class BinaryBox(Box):
    """A box of n binary values."""

    n: int

    def to(self, dest: torch.dtype | DEVICE_TYPING) -> ContinuousBox:
        return deepcopy(self)

    def __repr__(self):
        return f"{self.__class__.__name__}(n={self.n})"


@dataclass(repr=False)
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,
    or sometimes to simulate simple behaviors by generating random data within a defined space.

    TensorSpecs are primarily used in environments to specify their input/output structure without needing to
    execute the environment (or starting it). They can also be used to instantiate shared buffers to pass
    data from worker to worker.

    TensorSpecs are dataclasses that always share the following fields: `shape`, `space, `dtype` and `device`.

    As such, TensorSpecs possess some common behavior with :class:`~torch.Tensor` and :class:`~tensordict.TensorDict`:
    they can be reshaped, indexed, squeezed, unsqueezed, moved to another device etc.

    Args:
        shape (torch.Size): size of the tensor. The shape includes the batch dimensions as well as the feature
            dimension. A negative shape (``-1``) means that the dimension has a variable number of elements.
        space (Box): Box instance describing what kind of values can be expected.
        device (torch.device): device of the tensor.
        dtype (torch.dtype): dtype of the tensor.

    .. note:: A spec can be constructed from a :class:`~tensordict.TensorDict` using the :func:`~torchrl.envs.utils.make_composite_from_td`
        function. This function makes a low-assumption educated guess on the specs that may correspond to the input
        tensordict and can help to build specs automatically without an in-depth knowledge of the `TensorSpec` API.

    """

    shape: torch.Size
    space: None | Box
    device: torch.device | None = None
    dtype: torch.dtype = torch.float
    domain: str = ""
    _encode_memo_dict: dict[Any, Callable[[Any], Any]] = field(
        default_factory=dict,
    )

    SPEC_HANDLED_FUNCTIONS = {}

    def memoize_encode(self, mode: bool = True) -> None:
        """Creates a cached sequence of callables for the `encode` method that speeds up its execution.

        This should only be used whenever the input type, shape etc. are expected to be consistent across calls
        for a given spec.

        Args:
            mode (bool, optional): Whether the cache should be used. Defaults to `True`.

        .. seealso:: the cache can be erased via :meth:`~torchrl.data.TensorSpec.erase_memoize_cache`.
        """
        warnings.warn(
            "memoized encoding is an experimental feature. Use at your own risks."
        )
        if mode:
            self.encode = self._encode_memo
        else:
            self.encode = self._encode_eager

    def erase_memoize_cache(self) -> None:
        """Clears the memoized cache for cached encode execution.

        .. seealso:: :meth:`~torchrl.data.TensorSpec.memoize_encode`.
        """
        self._encode_memo_dict.clear()

    def __getstate__(self):
        state = dict(self.__dict__)
        state["_encode"] = {}
        return state

    @classmethod
    def implements_for_spec(cls, torch_function: Callable) -> Callable:
        """Register a torch function override for TensorSpec."""

        @wraps(torch_function)
        def decorator(func):
            cls.SPEC_HANDLED_FUNCTIONS[torch_function] = func
            return func

        return decorator

    @property
    def device(self) -> torch.device:
        """The device of the spec.

        Only :class:`Composite` specs can have a ``None`` device. All leaves must have a non-null device.
        """
        return self._device

    @device.setter
    def device(self, device: torch.device | None) -> None:
        self._device = _make_ordinal_device(device)

    def clear_device_(self) -> T:
        """A no-op for all leaf specs (which must have a device).

        For :class:`Composite` specs, this method will erase the device.
        """
        return self

    @abc.abstractmethod
    def cardinality(self) -> int:
        """The cardinality of the spec.

        This refers to the number of possible outcomes in a spec. It is assumed that the cardinality of a composite
        spec is the cartesian product of all possible outcomes.

        """
        ...

    def encode(
        self,
        val: np.ndarray | list | torch.Tensor | TensorDictBase,
        *,
        ignore_device: bool = False,
    ) -> torch.Tensor | TensorDictBase:
        """Encodes a value given the specified spec, and return the corresponding tensor.

        This method is to be used in environments that return a value (eg, a numpy array) that can be
        easily mapped to the TorchRL required domain.
        If the value is already a tensor, the spec will not change its value and return it as-is.

        Args:
            val (np.ndarray or torch.Tensor): value to be encoded as tensor.

        Keyword Args:
            ignore_device (bool, optional): if ``True``, the spec device will
                be ignored. This is used to group tensor casting within a call
                to ``TensorDict(..., device="cuda")`` which is faster.

        Returns:
            torch.Tensor matching the required tensor specs.

        """
        raise NotImplementedError(
            "This is a placeholder that needs to be set during construction"
        )

    def _encode_eager(
        self,
        val: np.ndarray | list | torch.Tensor | TensorDictBase,
        *,
        ignore_device: bool = False,
    ) -> torch.Tensor | TensorDictBase:
        if not isinstance(val, torch.Tensor):
            if isinstance(val, list):
                if len(val) == 1:
                    # gym used to return lists of images since 0.26.0
                    # We convert these lists in np.array or take the first element
                    # if there is just one.
                    # See https://github.com/pytorch/rl/pull/403/commits/73d77d033152c61d96126ccd10a2817fecd285a1
                    val = val[0]
                else:
                    val = np.array(val)
            if isinstance(val, np.ndarray) and not all(
                stride > 0 for stride in val.strides
            ):
                val = val.copy()
            if not ignore_device:
                val = torch.as_tensor(val, device=self.device, dtype=self.dtype)
            else:
                val = torch.as_tensor(val, dtype=self.dtype)
            if val.shape != self.shape:
                # if val.shape[-len(self.shape) :] != self.shape:
                # option 1: add a singleton dim at the end
                if val.shape == self.shape and self.shape[-1] == 1:
                    val = val.unsqueeze(-1)
                else:
                    try:
                        val = val.reshape(self.shape)
                    except Exception as err:
                        raise RuntimeError(
                            f"Shape mismatch: the value has shape {val.shape} which "
                            f"is incompatible with the spec shape {self.shape}."
                        ) from err
        if _CHECK_SPEC_ENCODE:
            self.assert_is_in(val)
        return val

    def _encode_memo(
        self,
        val: np.ndarray | list | torch.Tensor | TensorDictBase,
        *,
        ignore_device: bool = False,
    ) -> torch.Tensor | TensorDictBase:
        funcs = self._encode_memo_dict.get(ignore_device)
        if funcs is not None:
            return funcs(val)

        funcs = []
        val_orig = val
        if not isinstance(val, torch.Tensor):
            if isinstance(val, list):
                if len(val) == 1:
                    # gym used to return lists of images since 0.26.0
                    # We convert these lists in np.array or take the first element
                    # if there is just one.
                    # See https://github.com/pytorch/rl/pull/403/commits/73d77d033152c61d96126ccd10a2817fecd285a1
                    funcs.append(lambda val: val[0])
                else:
                    funcs.append(lambda val: np.array(val))
            val = _reduce_funcs(funcs)(val_orig)
            if isinstance(val, np.ndarray) and not all(
                stride > 0 for stride in val.strides
            ):
                funcs.append(lambda val: val.copy())
            val = _reduce_funcs(funcs)(val_orig)
            if not ignore_device:
                funcs.append(
                    lambda val: torch.as_tensor(
                        val, device=self.device, dtype=self.dtype
                    )
                )
            else:
                funcs.append(lambda val: torch.as_tensor(val, dtype=self.dtype))
            val = _reduce_funcs(funcs)(val_orig)
        if val.shape != self.shape:
            # if val.shape[-len(self.shape) :] != self.shape:
            # option 1: add a singleton dim at the end
            if val.shape == self.shape and self.shape[-1] == 1:
                funcs.append(lambda val: val.unsqueeze(-1))
            else:

                def reshape(val):
                    try:
                        return val.reshape(self.shape)
                    except Exception as err:
                        raise RuntimeError(
                            f"Shape mismatch: the value has shape {val.shape} which "
                            f"is incompatible with the spec shape {self.shape}."
                        ) from err

                funcs.append(reshape)
            val = _reduce_funcs(funcs)(val_orig)
        if _CHECK_SPEC_ENCODE:

            def check(val):
                self.assert_is_in(val)
                return val

            funcs.append(check)
        if len(funcs) == 0:
            self._encode_memo_dict[ignore_device] = lambda x: x
        elif len(funcs) == 1:
            self._encode_memo_dict[ignore_device] = funcs[0]
        else:
            self._encode_memo_dict[ignore_device] = _reduce_funcs(funcs)
        return self._encode_memo_dict[ignore_device](val_orig)

    @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)

    def __setattr__(self, key, value):
        if key == "shape":
            value = _size(value)
        super().__setattr__(key, value)

    def to_numpy(
        self, val: torch.Tensor | TensorDictBase, safe: bool = None
    ) -> np.ndarray | dict:
        """Returns the ``np.ndarray`` correspondent of an input tensor.

        This is intended to be the inverse operation of :meth:`.encode`.

        Args:
            val (torch.Tensor): tensor to be transformed_in to numpy.
            safe (bool): boolean value indicating whether a check should be
                performed on the value against the domain of the spec.
                Defaults to the value of the ``CHECK_SPEC_ENCODE`` environment variable.

        Returns:
            a np.ndarray.

        """
        if safe is None:
            safe = _CHECK_SPEC_ENCODE
        if safe:
            self.assert_is_in(val)
        return val.detach().cpu().numpy()

    @property
    def ndim(self) -> int:
        """Number of dimensions of the spec shape.

        Shortcut for ``len(spec.shape)``.

        """
        return self.ndimension()

    def ndimension(self) -> int:
        """Number of dimensions of the spec shape.

        Shortcut for ``len(spec.shape)``.

        """
        return len(self.shape)

    @property
    def _safe_shape(self) -> torch.Size:
        """Returns a shape where all heterogeneous values are replaced by one (to be expandable)."""
        return _size([int(v) if v >= 0 else 1 for v in self.shape])

    @abc.abstractmethod
    def index(
        self, index: INDEX_TYPING, tensor_to_index: torch.Tensor | TensorDictBase
    ) -> 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)

        """
        ...

    @overload
    def expand(self, shape: torch.Size):
        ...

    @abc.abstractmethod
    def expand(self, *shape: int) -> T:
        """Returns a new Spec with the expanded shape.

        Args:
            *shape (tuple or iterable of int): the new shape of the Spec.
                Must be broadcastable with the current shape:
                its length must be at least as long as the current shape length,
                and its last values must be compliant too; ie they can only differ
                from it if the current dimension is a singleton.

        """
        ...

    def squeeze(self, dim: int | None = None) -> T:
        """Returns a new Spec with all the dimensions of size ``1`` removed.

        When ``dim`` is given, a squeeze operation is done only in that dimension.

        Args:
            dim (int or None): the dimension to apply the squeeze operation to

        """
        shape = _squeezed_shape(self.shape, dim)
        if shape is None:
            return self
        return self.__class__(shape=shape, device=self.device, dtype=self.dtype)

    def unsqueeze(self, dim: int) -> T:
        """Returns a new Spec with one more singleton dimension (at the position indicated by ``dim``).

        Args:
            dim (int or None): the dimension to apply the unsqueeze operation to.

        """
        shape = _unsqueezed_shape(self.shape, dim)
        return self.__class__(shape=shape, device=self.device, dtype=self.dtype)

    def make_neg_dim(self, dim: int) -> T:
        """Converts a specific dimension to ``-1``."""
        if dim < 0:
            dim = self.ndim + dim
        if dim < 0 or dim > self.ndim - 1:
            raise ValueError(f"dim={dim} is out of bound for ndim={self.ndim}")
        self.shape = _size([s if i != dim else -1 for i, s in enumerate(self.shape)])

    @overload
    def reshape(self, shape) -> T:
        ...

    def reshape(self, *shape) -> T:
        """Reshapes a ``TensorSpec``.

        Check :func:`~torch.reshape` for more information on this method.

        """
        if len(shape) == 1 and not isinstance(shape[0], int):
            return self.reshape(*shape[0])
        return self._reshape(shape)

    view = reshape

    @abc.abstractmethod
    def _reshape(self, shape: torch.Size) -> T:
        ...

    def unflatten(self, dim: int, sizes: tuple[int]) -> T:
        """Unflattens a ``TensorSpec``.

        Check :func:`~torch.unflatten` for more information on this method.

        """
        return self._unflatten(dim, sizes)

    def _unflatten(self, dim: int, sizes: tuple[int]) -> T:
        shape = torch.zeros(self.shape, device="meta").unflatten(dim, sizes).shape
        return self._reshape(shape)

    def flatten(self, start_dim: int, end_dim: int) -> T:
        """Flattens a ``TensorSpec``.

        Check :func:`~torch.flatten` for more information on this method.

        """
        return self._flatten(start_dim, end_dim)

    def _flatten(self, start_dim, end_dim):
        shape = torch.zeros(self.shape, device="meta").flatten(start_dim, end_dim).shape
        return self._reshape(shape)

    @abc.abstractmethod
    def _project(
        self, val: torch.Tensor | TensorDictBase
    ) -> torch.Tensor | TensorDictBase:
        raise NotImplementedError(type(self))

    @abc.abstractmethod
    def is_in(self, val: torch.Tensor | TensorDictBase) -> bool:
        """If the value ``val`` could have been generated by the ``TensorSpec``, returns ``True``, otherwise ``False``.

        More precisely, the ``is_in`` methods checks that the value ``val`` is within the limits defined by the ``space``
        attribute (the box), and that the ``dtype``, ``device``, ``shape`` potentially other metadata match those
        of the spec. If any of these checks fails, the ``is_in`` method will return ``False``.

        Args:
            val (torch.Tensor): value to be checked.

        Returns:
            boolean indicating if values belongs to the TensorSpec box.

        """
        ...