class torch.Size, space: Union[None, Box], device: torch.device | None = None, dtype: torch.dtype = torch.float32, domain: str = '')[source]

Parent class of the tensor meta-data containers for observation, actions and rewards.

  • shape (torch.Size) – size of the tensor

  • 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

assert_is_in(value: Tensor) None[source]

Asserts whether a tensor belongs to the box, and raises an exception otherwise.


value (torch.Tensor) – value to be checked.


A no-op for all leaf specs (which must have a device).

encode(val: Union[ndarray, Tensor], *, ignore_device=False) Tensor[source]

Encodes a value given the specified spec, and return the corresponding tensor.


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

Keyword Arguments:

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.


torch.Tensor matching the required tensor specs.

abstract expand(*shape)[source]

Returns a new Spec with the extended shape.


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

flatten(start_dim, end_dim)[source]

Flattens a tensorspec.

Check flatten() for more information on this method.

classmethod implements_for_spec(torch_function: Callable) Callable[source]

Register a torch function override for TensorSpec.

abstract index(index: Union[int, Tensor, ndarray, slice, List], tensor_to_index: Tensor) Tensor[source]

Indexes the input tensor.

  • index (int, torch.Tensor, slice or list) – index of the tensor

  • tensor_to_index – tensor to be indexed


indexed tensor

abstract is_in(val: Tensor) bool[source]

If the value val is in the box defined by the TensorSpec, returns True, otherwise False.


val (torch.Tensor) – value to be checked


boolean indicating if values belongs to the TensorSpec box

project(val: Tensor) Tensor[source]

If the input tensor is not in the TensorSpec box, it maps it back to it given some heuristic.


val (torch.Tensor) – tensor to be mapped to the box.


a torch.Tensor belonging to the TensorSpec box.

abstract rand(shape=None) Tensor[source]

Returns a random tensor in the box. The sampling will be uniform unless the box is unbounded.


shape (torch.Size) – shape of the random tensor


a random tensor sampled in the TensorSpec box.


Reshapes a tensorspec.

Check reshape() for more information on this method.

squeeze(dim: int | None = None)[source]

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.


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

to_numpy(val: Tensor, safe: Optional[bool] = None) ndarray[source]

Returns the np.ndarray correspondent of an input tensor.

  • 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.


a np.ndarray

type_check(value: Tensor, key: Optional[str] = None) None[source]

Checks the input value dtype against the TensorSpec dtype and raises an exception if they don’t match.

  • value (torch.Tensor) – tensor whose dtype has to be checked

  • key (str, optional) – if the TensorSpec has keys, the value dtype will be checked against the spec pointed by the indicated key.

unflatten(dim, sizes)[source]

Unflattens a tensorspec.

Check unflatten() for more information on this method.


Reshapes a tensorspec.

Check reshape() for more information on this method.

zero(shape=None) Tensor[source]

Returns a zero-filled tensor in the box.


shape (torch.Size) – shape of the zero-tensor


a zero-filled tensor sampled in the TensorSpec box.


Access comprehensive developer documentation for PyTorch

View Docs


Get in-depth tutorials for beginners and advanced developers

View Tutorials


Find development resources and get your questions answered

View Resources