[Feature Request] Improvements to LazyStackedTensorSpec, LazyStackedCompositeSpec

Motivation

#892 introduced two new classes: LazyStackedTensorSpec, and LazyStackedCompositeSpec. These allow us to stack heterogeneous TensorSpec and CompositeSpec respectively.

These classes are missing a few features that still need to be added.

TODO

• Add documentation - the classes have basic docstrings, but we should add examples etc.
• Add support for nestedtensors
• #960 changed the behaviour of torch.stack to only return lazy stacks in the heterogenous case. As a result a number of tests now only cover the homogeneous or heterogeneous case but not both. These tests need to be expanded to cover both cases.

Unimplemented methods - the following exist as unimplemented placeholder methods in LazyStackedTensorSpec and LazyStackedCompositeSpec respectively

LazyStackedTensorSpec

• __eq__ (requires unbind?)
• __len__
• project
• is_in
• encode (not relevant, should be removed)
• __iter__
• __setitem__

LazyStackedCompositeSpec

• update
• __eq__
• __len__
• project
• is_in
• type_check
• encode
• __delitem__
• __iter__
• __setitem__