class_docstr: Message weighting by attention.
code: def forward(
      self,
      source: torch.LongTensor,
      target: torch.LongTensor,
      message: Optional[torch.FloatTensor] = None,
      x_e: Optional[torch.FloatTensor] = None,
  ) -> torch.FloatTensor:  # noqa: D102
      if message is None or x_e is None:
          raise ValueError(
              f"{self.__class__.__name__} requires message and x_e."
          )
  
      # view for heads
      message_ = message.view(message.shape[0], self.num_heads, -1)
      # compute attention coefficients, shape: (num_edges, num_heads)
      alpha = self.activation(
          einsum(
              "ihd,hd->ih",
              torch.cat(
                  [
                      message_,
                      x_e[target].view(
                          target.shape[0], self.num_heads, -1
                      ),
                  ],
                  dim=-1,
              ),
              self.weight,
          )
      )
      # TODO we can use scatter_softmax from torch_scatter directly, kept this if we can rewrite it w/o scatter
      alpha = softmax(
          alpha, index=target, num_nodes=x_e.shape[0], dim=0
      )
      alpha = self.dropout(alpha)
      return (message_ * alpha.view(-1, self.num_heads, 1)).view(
          -1, self.num_heads * self.attention_dim
      )
  
imports: from abc import abstractmethod
  from typing import ClassVar, Optional, Union
  import torch
  from class_resolver import ClassResolver
  from torch import nn
  from ..utils import einsum
label_desc: Use image node_id 2 for calling the AttentionEdgeWeighting obj's underlying member method code with example usage: obj.forward(source, target, message, x_e) and returns: unknown