Description

Fixes a bug introduced by #3077 where shapes for newly instantiated Binary specs are incompatible with the value of n passed.

In particular the following pattern is used in unsqueeze and other shape-modifying operations:

def unsqueeze(self, dim: int):
        shape = _unsqueezed_shape(self.shape, dim)
        return self.__class__(
            n=self.shape[-1] if len(self.shape) > 0 else None,
            shape=shape,
            device=self.device,
            dtype=self.dtype,
        )

However, this is wrong because we should extract n from the new shape rather than from the current self.shape.

We fix by switching to check on the new shape:

def unsqueeze(self, dim: int):
        shape = _unsqueezed_shape(self.shape, dim)
        return self.__class__(
            n=shape[-1] if len(shape) > 0 else None,
            shape=shape,
            device=self.device,
            dtype=self.dtype,
        )

Motivation and Context

Without this change, we have value errors raised due to shape mismatches:

import pytest
import torch
from torchrl.data import Binary


def test__Binary__if_empty_cannot_expand() -> None:
    spec = Binary(shape=())
    new_shape = (3, 3)

    with pytest.raises(ValueError, match="The last value of the shape must match 'n'"):
        spec.expand(new_shape)

NOTE: I have tried to include a test to check for this, but I don't know if that is the correct place and scope.

Types of changes

What types of changes does your code introduce? Remove all that do not apply:

• Bug fix (non-breaking change which fixes an issue)
• New feature (non-breaking change which adds core functionality)
• Breaking change (fix or feature that would cause existing functionality to change)
• Documentation (update in the documentation)
• Example (update in the folder of examples)

Checklist

Go over all the following points, and put an x in all the boxes that apply.
If you are unsure about any of these, don't hesitate to ask. We are here to help!

• I have read the CONTRIBUTION guide
• My change requires a change to the documentation.
• I have updated the tests accordingly (required for a bug fix or a new feature).
• I have updated the documentation accordingly.