How to store non-aligned data?

[Roy-Kid]

Hi everyone,

I try to deal with the chemical dataset, which has a lot of small molecules with different atoms. How to use torchdict to store those non-aligned data, for example, coordinates. What should I set the batch_size?

[vmoens]

Good question
Can you give an example of the data you want to bring together and what you want to do with it?

One thing I'm interested in is to provide a nice solution to build TDs with bin evenly spaced data that come in stream (say stream 1 gives you data at 7Hz and stream 2 gives you data at 3Hz), and you want to build TD that can be indexed along the time dimension at, say, 5Hz.
One way to do it is transition data through buffers and dump the data in separate TDs that you stack using LazyStackedTD.
Another (or complementary) solution is to bring the data together continuously with nested tensors.

Is your problem similar?
Eg you have N molecules that you want to batch in a single TD along the shape dimension, but your molecules have a different shape?

[Roy-Kid]

Hi, I tried it, and it looks nice. I found it can not use key-select, for example:

mol1 = {
    'xyz': torch.randn(2, 3),
    'energy': torch.tensor([1.0]),
}
mol2 = {
    'xyz': torch.randn(4, 3),
    'energy': torch.tensor([2.0]),
}
mol3 = {
    'xyz': torch.randn(5, 3),
    'energy': torch.tensor([3.0]),
}
td1 = TensorDict(mol1)
td2 = TensorDict(mol2)
td3 = TensorDict(mol3)

tds = TensorDict.lazy_stack([td1, td2, td3])

# ok!
tds[0]

# not ok
tds.get('xyz')
# untimeError: Found more than one unique shape in the tensors to be stacked ({torch.Size([2, 3]), torch.Size([5, 3]), torch.Size([4, 3])}). This is likely due to a modification of one of the stacked TensorDicts, where a key has been updated/created with an uncompatible shape. If the entries are intended to have a different shape, use the get_nestedtensor method instead.

# but I can do:
coords = tds.get_nestedtensor('xyz')
atom_number_list = torch.tensor([len(xyz) for xyz in coords])
# tensor([2, 4, 5])
batch_id = torch.repeat_interleave(atom_number_list)
# tensor([0, 0, 1, 1, 1, 1, 2, 2, 2, 2, 2])

Is this a right and elegant way to use it? Because I think when fed lazystacktensordict into a module, the module itself still needs to concatenate nested tensordict.

[vmoens]

I have a PR open to consolidate all ragged shapes in a nested tensor
#785
Something like

td.get("xyz") # Error: wrong shape
td = td.densify()
td.get("xyz") # Gives a nested tensor

Would that help?

[Roy-Kid]

Super! it is very useful!

So I generate mask manually is a good way right? It should be almost the elegant way to do that

[vmoens]

#955
Merged!
If you use nightlies (from tomorrow) you will be able to call densify and all the entries will be accessible.
From there, you can call to_padded_tensor if you need to make these nested tensors "regular" ones.
If needed I can make another PR that adds a padded option in densify.

[Roy-Kid]

super! I have already started to refactor the code. It is more promising than my pathetic dict! Thanks a lot for help

[Roy-Kid]

Brilliant! Very straightforward to use!

I have a question. Suppose I have a sequence of nested TensorDict

raw_frames = []
for n_atoms in [6, 8, 10]:

    frame = TensorDict({
        'atoms': {
            'xyz': torch.randn(n_atoms, 3)
        },
        'pairs': {
            'pair_i': torch.randint(0, n_atoms, (n_atoms*3, 2)),
            'pair_j': torch.randint(0, n_atoms, (n_atoms*3, 2))
        },


    })

    raw_frames.append(frame)

raw_frames = TensorDict.maybe_dense_stack(raw_frames)
dense_frames = frames.densify()

I want to get padded xyz:

dense_frames['atoms'].to_padded_tensor(0.0, 'atoms_mask')  # error
# raise error:
# RuntimeError: mask_key should only be provided if the heterogenous dimension is part of the batch-size.
# which dense_frames['atoms'] is:
# TensorDict(
#    fields={
#        xyz: Tensor(shape=torch.Size([3, -1, 3]), device=cpu, dtype=torch.float32, is_shared=False)},
#    batch_size=torch.Size([3]),
#    device=None,
#    is_shared=False)
# and dense_frames['atoms', 'xyz'] is:
# nested_tensor([tensor1, tensor2, tensor3])

I am confused as to why it does not match. The heterogenous dimension should refer to -1, right? Maybe I need to .reshape(-1, 3) -> shape(-1, 3) manually?

Also, is there any API to get the mask if I don't want to padding?

atom_number_list = torch.tensor([len(xyz) for xyz in coords])
# tensor([2, 4, 5])
batch_id = torch.repeat_interleave(atom_number_list)
# tensor([0, 0, 1, 1, 1, 1, 2, 2, 2, 2, 2])

[vmoens]

Regarding the first point: I'm AFK atm but my guess is that it should work if you work with ["atoms", "xyz"] directly:

https://pytorch.org/docs/stable/nested.html#torch.nested.to_padded_tensor

[vmoens]

Oh you could also say that each tensordict "atoms" has shape "n_atoms"!
That will make the function accessible