EdgeWeighting

class EdgeWeighting(**kwargs)[source]

Bases: torch.nn.modules.module.Module

Base class for edge weightings.

Initializes internal Module state, shared by both nn.Module and ScriptModule.

Attributes Summary

needs_message

whether the edge weighting needs access to the message

Methods Summary

forward(source, target[, message, x_e])

Compute edge weights.

Attributes Documentation

needs_message: ClassVar[bool] = False

whether the edge weighting needs access to the message

Methods Documentation

abstract forward(source, target, message=None, x_e=None)[source]

Compute edge weights.

Parameters

  • source (LongTensor) – shape: (num_edges,) The source indices.

  • target (LongTensor) – shape: (num_edges,) The target indices.

  • message (Optional[FloatTensor]) – shape (num_edges, dim) Actual messages to weight

  • x_e (Optional[FloatTensor]) – shape (num_nodes, dim) Node states up to the weighting point

Return type

FloatTensor

Returns

shape: (num_edges, dim) Messages weighted with the edge weights.