easygraph.nn.convs.pma module

class easygraph.nn.convs.pma.PMA(in_channels, hid_dim, out_channels, num_layers, heads=1, concat=True, negative_slope=0.2, dropout=0.0, bias=False, **kwargs)

Bases: MessagePassing

PMA part: Note that in original PMA, we need to compute the inner product of the seed and neighbor nodes. i.e. e_ij = a(Wh_i,Wh_j), where a should be the inner product, h_i is the seed and h_j are neightbor nodes. In GAT, a(x,y) = a^T[x||y]. We use the same logic.

Attributes:

decomposed_layers

explain

Methods

add_module(name, module)	Adds a child module to the current module.
aggregate(inputs, index[, dim_size, aggr])	Aggregates messages from neighbors as \(\square_{j \in \mathcal{N}(i)}\).
apply(fn)	Applies fn recursively to every submodule (as returned by .children()) as well as self.
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.
edge_update()	Computes or updates features for each edge in the graph.
edge_updater(edge_index[, size])	The initial call to compute or update features for each edge in the graph.
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, edge_index[, size, ...])	param return_attention_weights: If set to True,
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.
jittable([typing])	Analyzes the MessagePassing instance and produces a new jittable module that can be used in combination with torch.jit.script().
load_state_dict(state_dict[, strict])	Copies parameters and buffers from state_dict into this module and its descendants.
message(x_j, alpha_j, index, ptr, size_j)	Constructs messages from node \(j\) to node \(i\) in analogy to \(\phi_{\mathbf{\Theta}}\) for each edge in edge_index.
message_and_aggregate(edge_index)	Fuses computations of message() and aggregate() into a single function.
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.
propagate(edge_index[, size])	The initial call to start propagating messages.
register_aggregate_forward_hook(hook)	Registers a forward hook on the module.
register_aggregate_forward_pre_hook(hook)	Registers a forward pre-hook on the module.
register_backward_hook(hook)	Registers a backward hook on the module.
register_buffer(name, tensor[, persistent])	Adds a buffer to the module.
register_edge_update_forward_hook(hook)	Registers a forward hook on the module.
register_edge_update_forward_pre_hook(hook)	Registers a forward pre-hook on 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_message_and_aggregate_forward_hook(hook)	Registers a forward hook on the module.
register_message_and_aggregate_forward_pre_hook(hook)	Registers a forward pre-hook on the module.
register_message_forward_hook(hook)	Registers a forward hook on the module.
register_message_forward_pre_hook(hook)	Registers a forward pre-hook on the module.
register_module(name, module)	Alias for add_module().
register_parameter(name, param)	Adds a parameter to the module.
register_propagate_forward_hook(hook)	Registers a forward hook on the module.
register_propagate_forward_pre_hook(hook)	Registers a forward pre-hook on 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()	Resets all learnable parameters of the module.
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(inputs)	Updates node embeddings in analogy to \(\gamma_{\mathbf{\Theta}}\) for each node \(i \in \mathcal{V}\).
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__
explain_message

aggregate(inputs, index, dim_size=None, aggr='add')

Aggregates messages from neighbors as \(\square_{j \in \mathcal{N}(i)}\).

Takes in the output of message computation as first argument and any argument which was initially passed to propagate().

By default, this function will delegate its call to scatter functions that support “add”, “mean” and “max” operations as specified in __init__() by the aggr argument.

forward(x, edge_index: Tensor | SparseTensor, size: Tuple[int, int] | None = None, return_attention_weights=None)

Parameters:

return_attention_weights (bool, optional) – If set to True, will additionally return the tuple (edge_index, attention_weights), holding the computed attention weights for each edge. (default: None)

message(x_j, alpha_j, index, ptr, size_j)

Constructs messages from node \(j\) to node \(i\) in analogy to \(\phi_{\mathbf{\Theta}}\) for each edge in edge_index. This function can take any argument as input which was initially passed to propagate(). Furthermore, tensors passed to propagate() can be mapped to the respective nodes \(i\) and \(j\) by appending _i or _j to the variable name, .e.g. x_i and x_j.

reset_parameters()

Resets all learnable parameters of the module.

easygraph.nn.convs.pma.glorot(tensor)

easygraph.nn.convs.pma.zeros(tensor)