Message Passing

Message passing module.

class topomodelx.base.message_passing.MessagePassing(aggr_func: Literal['sum', 'mean', 'add'] = 'sum', att: bool = False, initialization: Literal['uniform', 'xavier_uniform', 'xavier_normal'] = 'xavier_uniform', initialization_gain: float = 1.414)

Define message passing.

This class defines message passing through a single neighborhood N, by decomposing it into 2 steps:

1. 🟥 Create messages going from source cells to target cells through N.

2. 🟧 Aggregate messages coming from different sources cells onto each target cell.

This class should not be instantiated directly, but rather inherited through subclasses that effectively define a message passing function.

This class does not have trainable weights, but its subclasses should define these weights.

Parameters:

aggr_funcLiteral[“sum”, “mean”, “add”], default=”sum”

Aggregation function to use.

attbool, default=False

Whether to use attention.

initializationLiteral[“uniform”, “xavier_uniform”, “xavier_normal”], default=”xavier_uniform”

Initialization method for the weights of the layer.

initialization_gainfloat, default=1.414

Gain for the weight initialization.

Methods

add_module(name, module)	Adds a child module to the current module.
aggregate(x_message)	Aggregate messages on each target cell.
apply(fn)	Applies fn recursively to every submodule (as returned by .children()) as well as self.
attention(x_source[, x_target])	Compute attention weights for messages.
bfloat16()	Casts all floating point parameters and buffers to bfloat16 datatype.
buffers([recurse])	Returns an iterator over module buffers.
children()	Returns an iterator over immediate children modules.
cpu()	Moves all model parameters and buffers to the CPU.
cuda([device])	Moves all model parameters and buffers to the GPU.
double()	Casts all floating point parameters and buffers to double datatype.
eval()	Sets the module in evaluation mode.
extra_repr()	Set the extra representation of the module
float()	Casts all floating point parameters and buffers to float datatype.
forward(x_source, neighborhood[, x_target])	Forward pass.
get_buffer(target)	Returns the buffer given by target if it exists, otherwise throws an error.
get_extra_state()	Returns any extra state to include in the module's state_dict.
get_parameter(target)	Returns the parameter given by target if it exists, otherwise throws an error.
get_submodule(target)	Returns the submodule given by target if it exists, otherwise throws an error.
half()	Casts all floating point parameters and buffers to half datatype.
ipu([device])	Moves all model parameters and buffers to the IPU.
load_state_dict(state_dict[, strict])	Copies parameters and buffers from state_dict into this module and its descendants.
message(x_source[, x_target])	Construct message from source cells to target cells.
modules()	Returns an iterator over all modules in the network.
named_buffers([prefix, recurse, ...])	Returns an iterator over module buffers, yielding both the name of the buffer as well as the buffer itself.
named_children()	Returns an iterator over immediate children modules, yielding both the name of the module as well as the module itself.
named_modules([memo, prefix, remove_duplicate])	Returns an iterator over all modules in the network, yielding both the name of the module as well as the module itself.
named_parameters([prefix, recurse, ...])	Returns an iterator over module parameters, yielding both the name of the parameter as well as the parameter itself.
parameters([recurse])	Returns an iterator over module parameters.
register_backward_hook(hook)	Registers a backward hook on the module.
register_buffer(name, tensor[, persistent])	Adds a buffer to the module.
register_forward_hook(hook, *[, prepend, ...])	Registers a forward hook on the module.
register_forward_pre_hook(hook, *[, ...])	Registers a forward pre-hook on the module.
register_full_backward_hook(hook[, prepend])	Registers a backward hook on the module.
register_full_backward_pre_hook(hook[, prepend])	Registers a backward pre-hook on the module.
register_load_state_dict_post_hook(hook)	Registers a post hook to be run after module's load_state_dict is called.
register_module(name, module)	Alias for add_module().
register_parameter(name, param)	Adds a parameter to the module.
register_state_dict_pre_hook(hook)	These hooks will be called with arguments: self, prefix, and keep_vars before calling state_dict on self.
requires_grad_([requires_grad])	Change if autograd should record operations on parameters in this module.
reset_parameters()	Reset learnable parameters.
set_extra_state(state)	This function is called from load_state_dict() to handle any extra state found within the state_dict.
share_memory()	See torch.Tensor.share_memory_()
state_dict(*args[, destination, prefix, ...])	Returns a dictionary containing references to the whole state of the module.
to(*args, **kwargs)	Moves and/or casts the parameters and buffers.
to_empty(*, device)	Moves the parameters and buffers to the specified device without copying storage.
train([mode])	Sets the module in training mode.
type(dst_type)	Casts all parameters and buffers to dst_type.
xpu([device])	Moves all model parameters and buffers to the XPU.
zero_grad([set_to_none])	Sets gradients of all model parameters to zero.

__call__

References

[1]

Hajij, Zamzmi, Papamarkou, Miolane, Guzmán-Sáenz, Ramamurthy, Birdal, Dey, Mukherjee, Samaga, Livesay, Walters, Rosen, Schaub. Topological deep learning: going beyond graph data (2023). https://arxiv.org/abs/2206.00606.

[2]

Papillon, Sanborn, Hajij, Miolane. Architectures of topological deep learning: a survey on topological neural networks (2023). https://arxiv.org/abs/2304.10031.

aggregate(x_message)

Aggregate messages on each target cell.

A target cell receives messages from several source cells. This function aggregates these messages into a single output feature per target cell.

🟧 This function corresponds to the within-neighborhood aggregation defined in [1]_ and [2]_.

Parameters:

x_messagetorch.Tensor, shape = (…, n_messages, out_channels)

Features associated with each message. One message is sent from a source cell to a target cell.

Returns:

Tensor, shape = (…, n_target_cells, out_channels)

Output features on target cells. Each target cell aggregates messages from several source cells. Assumes that all target cells have the same rank s.

attention(x_source, x_target=None)

Compute attention weights for messages.

This provides a default attention function to the message-passing scheme.

Alternatively, users can subclass MessagePassing and overwrite the attention method in order to replace it with their own attention mechanism.

The implementation follows [1]_.

Parameters:

x_sourcetorch.Tensor, shape = (n_source_cells, in_channels)

Input features on source cells. Assumes that all source cells have the same rank r.

x_targettorch.Tensor, shape = (n_target_cells, in_channels)

Input features on source cells. Assumes that all source cells have the same rank r.

Returns:

torch.Tensor, shape = (n_messages, 1)

Attention weights: one scalar per message between a source and a target cell.

forward(x_source, neighborhood, x_target=None)

Forward pass.

This implements message passing for a given neighborhood:

• from source cells with input features x_source,

• via neighborhood defining where messages can pass,

• to target cells with input features x_target.

In practice, this will update the features on the target cells.

If not provided, x_target is assumed to be x_source, i.e. source cells send messages to themselves.

The message passing is decomposed into two steps:

1. 🟥 Message: A message \(m_{y \rightarrow x}^{\left(r \rightarrow s\right)}\) travels from a source cell \(y\) of rank r to a target cell \(x\) of rank s through a neighborhood of \(x\), denoted \(\mathcal{N} (x)\), via the message function \(M_\mathcal{N}\):

\[m_{y \rightarrow x}^{\left(r \rightarrow s\right)} = M_{\mathcal{N}}\left(\mathbf{h}_x^{(s)}, \mathbf{h}_y^{(r)}, \Theta \right),\]

where:

• \(\mathbf{h}_y^{(r)}\) are input features on the source cells, called x_source,

• \(\mathbf{h}_x^{(s)}\) are input features on the target cells, called x_target,

• \(\Theta\) are optional parameters (weights) of the message passing function.

Optionally, attention can be applied to the message, such that:

\[m_{y \rightarrow x}^{\left(r \rightarrow s\right)} \leftarrow att(\mathbf{h}_y^{(r)}, \mathbf{h}_x^{(s)}) . m_{y \rightarrow x}^{\left(r \rightarrow s\right)}\]

2. 🟧 Aggregation: Messages are aggregated across source cells \(y\) belonging to the neighborhood \(\mathcal{N}(x)\):

\[m_x^{\left(r \rightarrow s\right)} = \text{AGG}_{y \in \mathcal{N}(x)} m_{y \rightarrow x}^{\left(r\rightarrow s\right)},\]

resulting in the within-neighborhood aggregated message \(m_x^{\left(r \rightarrow s\right)}\).

Details can be found in [1]_ and [2]_.

Parameters:

x_sourceTensor, shape = (…, n_source_cells, in_channels)

Input features on source cells. Assumes that all source cells have the same rank r.

neighborhoodtorch.sparse, shape = (n_target_cells, n_source_cells)

Neighborhood matrix.

x_targetTensor, shape = (…, n_target_cells, in_channels)

Input features on target cells. Assumes that all target cells have the same rank s. Optional. If not provided, x_target is assumed to be x_source, i.e. source cells send messages to themselves.

Returns:

torch.Tensor, shape = (…, n_target_cells, out_channels)

Output features on target cells. Assumes that all target cells have the same rank s.

message(x_source, x_target=None)

Construct message from source cells to target cells.

🟥 This provides a default message function to the message passing scheme.

Alternatively, users can subclass MessagePassing and overwrite the message method in order to replace it with their own message mechanism.

Parameters:

x_sourceTensor, shape = (…, n_source_cells, in_channels)

Input features on source cells. Assumes that all source cells have the same rank r.

x_targetTensor, shape = (…, n_target_cells, in_channels)

Input features on target cells. Assumes that all target cells have the same rank s. Optional. If not provided, x_target is assumed to be x_source, i.e. source cells send messages to themselves.

Returns:

torch.Tensor, shape = (…, n_source_cells, in_channels)

Messages on source cells.

reset_parameters()

Reset learnable parameters.

Notes

This function will be called by subclasses of MessagePassing that have trainable weights.