lloca.backbone.lloca_message_passing.LLoCaMessagePassing

class lloca.backbone.lloca_message_passing.LLoCaMessagePassing(params_dict, aggr='add')[source]

Bases: MessagePassing

Adaptation of the torch_geometric MessagePassing class using the LLoCa formalism.

Initializes a new instance of the LLoCaMessagePassing class.

Parameters:
  • params_dict (dict[str, dict[str, Any]]) – A dictionary containing the parameters for the message passing algorithm and the corresponding representations. Each key in the dictionary represents a feature, and the value is another dictionary with keys “type” and “rep”. The “type” can be either “local” or “global”, and “rep” is an instance of TensorRepsTransform that defines how to transform the features.

  • aggr (str, optional) – The aggregation method to use for combining messages. Defaults to “add”.

pre_message_hook(module, inputs)[source]

Pre-message hook method that is called before passing messages in the message passing algorithm. We transform the features according to the representations in the params_dict.

Return type:

tuple

pre_propagate_hook(module, inputs)[source]

A hook method called before propagating messages in the message passing algorithm. We save the frames in the class variable and remove it from the inputs dictionary.

Return type:

tuple