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
, andkeep_vars
before callingstate_dict
onself
.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.