[Feature] dense_stack_tds
#506
Conversation
facebook-github-bot
added
the
CLA Signed
tensordict/tensordict.py
Outdated
| This must be used when some of the :class:`tensordict.TensorDictBase` objects that need to be stacked | ||
| can have :class:`tensordict.LazyStackedTensorDict` among entries (or nested entries). |
There was a problem hiding this comment.
why "can have"? why not just "have"?
There was a problem hiding this comment.
because this method can be used even when no lazy stacks are involved
tensordict/tensordict.py
Outdated
| 1 -> | ||
| b: Tensor(shape=torch.Size([2, 2]), device=cpu, dtype=torch.float32, is_shared=False)}, | ||
| batch_size=torch.Size([2, 2]), | ||
| device=None, | ||
| is_shared=False, | ||
| stack_dim=1) |
There was a problem hiding this comment.
the appearance of b seems to be the only difference but it's still a lazy key
Why do we need 2 levels of nesting? This example isn't super clear to me
There was a problem hiding this comment.
we need 2 levels of nesting because the function is used to stack tds that contain lazy stacks. so it is only useful when there are at least two levels of nesting. The change of the stack dim from 0 to 1 is the core difference here. "b" is just there to highlight that
| "nested_stack_dim", | ||
| [0, 1, 2], | ||
| ) | ||
| def test_dense_stack_tds(stack_dim, nested_stack_dim): |
There was a problem hiding this comment.
I don't see what we're testing specific to this method here. All the assertions would hold with a regular tensordict no?
Why aren't we testing anything specific to the keys we provide?
There was a problem hiding this comment.
the idea is that this function should behave exactly the same with regualr tds or lazy stacks so the fact that a regualr td passes is ok to me. the main thing we are testing is that, when stacking tds that contain lazy stacks, we can call this function to remove one level of lazyness and densify one stack_dim
|
I have updated docstring, example and test. |
No description provided.