Conv#

Convolutional layer for message passing.

class topomodelx.base.conv.Conv(in_channels, out_channels, aggr_norm: bool = False, update_func: Literal['relu', 'sigmoid', None] = None, att: bool = False, initialization: Literal['xavier_uniform', 'xavier_normal'] = 'xavier_uniform', initialization_gain: float = 1.414, with_linear_transform: bool = True)[source]#

Message passing: steps 1, 2, and 3.

Builds the message passing route given by one neighborhood matrix. Includes an option for an x-specific update function.

Parameters:
in_channelsint

Dimension of input features.

out_channelsint

Dimension of output features.

aggr_normbool, default=False

Whether to normalize the aggregated message by the neighborhood size.

update_func{“relu”, “sigmoid”}, optional

Update method to apply to message.

attbool, default=False

Whether to use attention.

initialization{“xavier_uniform”, “xavier_normal”}, default=”xavier_uniform”

Initialization method.

initialization_gainfloat, default=1.414

Initialization gain.

with_linear_transformbool, default=True

Whether to apply a learnable linear transform. NB: if False in_channels has to be equal to out_channels.

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.

update(x_message_on_target)

Update embeddings on each cell (step 4).

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__

forward(x_source, neighborhood, x_target=None) Tensor[source]#

Forward pass.

This implements message passing: - 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.

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.

update(x_message_on_target) Tensor[source]#

Update embeddings on each cell (step 4).

Parameters:
x_message_on_targettorch.Tensor, shape = (n_target_cells, out_channels)

Output features on target cells.

Returns:
torch.Tensor, shape = (n_target_cells, out_channels)

Updated output features on target cells.