Does pyg support adding new attributes to Batch objects?

[genesis-jamin]

We're doing something like this in our code right now:

# batch is a HeteroDataBatch
batch['mean_pos'].pos = scatter_mean(batch['ligand'].pos, batch['ligand'].batch, 0)

However we find that later when we call batch.to_data_list(), we get an error like this:

  File "...", line 761, in inference
    predictions_list.extend(prediction.cpu().to_data_list())
  File "/home/jamin/.conda/envs/pt2/lib/python3.10/site-packages/torch_geometric/data/batch.py", line 193, in to_data_list
    return [self.get_example(i) for i in range(self.num_graphs)]
  File "/home/jamin/.conda/envs/pt2/lib/python3.10/site-packages/torch_geometric/data/batch.py", line 193, in <listcomp>
    return [self.get_example(i) for i in range(self.num_graphs)]
  File "/home/jamin/.conda/envs/pt2/lib/python3.10/site-packages/torch_geometric/data/batch.py", line 124, in get_example
    data = separate(
  File "/home/jamin/.conda/envs/pt2/lib/python3.10/site-packages/torch_geometric/data/separate.py", line 35, in separate
    attrs = slice_dict[key].keys()
KeyError: 'mean_pos'

What's the correct way to add new attributes to the Batch object?

[rusty1s]

In this case, you would also need to add mean_pos to _slice_dict, e.g.:

batch['mean_pos'].pos = ...
batch._slice_dict['mean_pos']['pos'] = torch.arange(0, batch['ligand'].max())

There exists some plans to make this easier, see #8414