tfg.geometry.convolution.graph_pooling.pool

Implements graph pooling.

tfg.geometry.convolution.graph_pooling.pool(
    data: type_alias.TensorLike,
    pool_map: type_alias.TensorLike,
    sizes: type_alias.TensorLike,
    algorithm: str = 'max',
    name: str = 'graph_pooling_pool'
) -> tf.Tensor

The features at each output vertex are computed by pooling over a subset of vertices in the input graph. This pooling window is specified by the input pool_map.

The shorthands used below are V1: The number of vertices in the input data. V2: The number of vertices in the pooled output data. C: The number of channels in the data.

In the following, A1 to An are optional batch dimensions.

data A float tensor with shape [A1, ..., An, V1, C].
pool_map A SparseTensor with the same type as data and with shape [A1, ..., An, V2, V1]. The features for an output vertex v2 will be computed by pooling over the corresponding input vertices specified by the entries in pool_map[A1, ..., An, v2, :].
sizes An int tensor of shape [A1, ..., An, 2] indicating the true input sizes in case of padding (sizes=None indicates no padding). sizes[A1, ..., An, 0] <= V2 specifies the padding in the (pooled) output, and sizes[A1, ..., An, 1] <= V1 specifies the padding in the input.
algorithm The pooling function, must be either 'max' or 'weighted'. Default is 'max'. For 'max' pooling, the output features are the maximum over the input vertices (in this case only the indices of the SparseTensor pool_map are used, the values are ignored). For 'weighted', the output features are a weighted sum of the input vertices, the weights specified by the values of pool_map.
name A name for this op. Defaults to 'graph_pooling_pool'.

Tensor with shape [A1, ..., An, V2, C].

TypeError if the input types are invalid.
ValueError if the input dimensions are invalid.
ValueError if algorithm is invalid.