About ptr in batches

[Nokimann]

When a batch of graphs is formed, we have additional information: batch and ptr

Global Pooling Layers commonly uses with batch

But where to use ptr?

Can someone show any examples?

[rusty1s]

We recently also added support in torch_scatter for segment reductions, i.e., reductions in which the indices are ordered and grouped together. As such, one alternative to implement global pooling layers can be achieved via:

from torch_scatter import segment_csr

out = segment_csr(x, data.ptr, reduce='mean')

which is equivalent to doing:

out = global_mean_pool(x, data.batch)

The main advantage is that segment_csr is deterministic, i.e., no atomic operations are used.

[Nokimann]

Thank @rusty1s
Does that mean segment_csr is a little faster operator than global_mean_pool?
It sounds like segment_csr does not pool nodes so that indices of the nodes match in data.batch.
Instead segment_csr just uses the range information in ptr.

[rusty1s]

segment_csr is slightly faster, but far from noticeable in practice. Note that the output of the two operations is equal, it's just that segment_csr has some constraints on how the input looks like (sorted groups denoted by range/boundary information in ptr).

[Nokimann]

@rusty1s
In segment_csr, the reduce argument only supports "sum", "mean", "min" or "max."

Does torch_scatter provide a concatenation function by segments?
Or, can we customize something like segment_cat easily?

[rusty1s]

What should segment_cat do here? I think this would be similar to utils.to_dense_batch, right?

[Nokimann]

This is what I expected! Thank you.

It seems not working in another discussion
But, I found a different way.