[Draft] Coding trees with TensorDict

[tcbegley]

Description

This PR contains a draft implementation of support for trees using TensorDict.

We allow the user to create a tree structure where each node in the tree is a tensordict with the same keys. The entire contents of the tree are backed by a single source tensordict with pre-allocated memory. As nodes are added to the tree the data in the node is stored in the source. The primary advantage of this setup is that collecting the data from multiple nodes requires us to simply index the source which can be done very efficiently compared to a naive alternative which would require us to iterate over the nodes we want and stack the results.

The interface is very preliminary and up for debate. The tests should give an idea of usage, here's a very basic example.

>>> import torch
>>> from tensordict import TensorDict
>>> from tensordict.prototype import make_tree

>>> # create the root node from a tensordict
>>> # preallocate memory for 100 nodes in the source
>>> root = make_tree(TensorDict({"data": torch.rand(2, 3, 4)}, [2, 3]), n_nodes=100)

>>> # adding a new tensordict as a child creates a new node in the tree
>>> root["left"] = TensorDict({"data": torch.rand(2, 3, 4)}, [2, 3])

>>> # we can now gather the data for the two nodes in the tree
>>> root.get_multiple_items("data", ("left", "data"))
tensor(...)  # shape torch.Size([2, 2, 3, 4])

>>> # alternatively we can create a tree from a nested tensordict
>>> root = make_tree(
...     TensorDict(
...         {"data": torch.ones(2, 3, 4), "left": TensorDict({"data": torch.zeros(2, 3, 4)}, [2, 3])},
...         [2, 3]
...     ),
...     n_nodes=100,
... )
>>> root.get_multiple_items("data", ("left", "data"))

[vmoens]

This is super cool!
Can we add a file in the benchmark directory to test the tree against a plain tensordict implementation (if that makes sense)?

[tcbegley]

Added some simple benchmarks. lmk what other things you think we should benchmark

get_multiple_items
tensordict: 2.145204999993439
tree: 0.9129869579919614

get_multiple_items_deep
tensordict 2.4592998750013066
tree 0.8956774580001365

[vmoens]

impressive!