Allow direct contruction of `Batch`

🚀 The feature, motivation and pitch

I’m working on GANs for graphs that generates multiple graphs at once. The output of the generator is a feature matrix and a vector in the form of Batch.batch. I would like to pass it on to the discriminator directly as a Batch, without splitting the matrix first, creating Data objects and then calling from_data_list. After the construction .__getitem__ and .num_graphs should be available.

Cheers!

Alternatives

No response

Additional context

No response

[rusty1s]

Glad that it’s working.

I think it is not really possible to allow this in a generic case, as we usually cannot utilize any assumptions on _slice_dict and _inc_dict layouts. How they look like is really task-dependent.

[mova]

Hey, thanks for your help. The following code works for now.

def batch_construct_direct(pcs: torch.Tensor, batch_idxs: torch.Tensor) -> Batch:
    device = pcs.device
    batch = Batch(x=pcs, batch=batch_idxs)
    batch._num_graphs = int(batch.batch.max() + 1)

    _slice_dict = defaultdict(dict)
    out = torch_scatter.scatter_add(
        torch.ones(len(batch.x), dtype=torch.long, device=device), batch_idxs, dim=0
    )
    out = out.cumsum(dim=0)
    _slice_dict["x"] = torch.cat(
        [torch.zeros(1, dtype=torch.long, device=device), out], dim=0
    )

    batch._slice_dict = _slice_dict

    _inc_dict = defaultdict(dict)
    _inc_dict["x"] = torch.zeros(batch._num_graphs, dtype=torch.long, device=device)
    batch._inc_dict = _inc_dict
    return batch
    ```
Do you think it would be worth a having a constructor of the `Batch` class could take care of this?  If not feel free to close this.