class_docstr: A representation with message passing with uses categorical relation type information.
  
  The message passing layers of this module internally handle the categorical relation type information
  via an `edge_type` input, e.g., :class:`torch_geometric.nn.conv.RGCNConv`, or
  :class:`torch_geometric.nn.conv.RGATConv`.
  
  The following example creates a one-layer RGCN using the basis decomposition:
  
  .. code-block:: python
  
      from pykeen.datasets import get_dataset
  
      embedding_dim = 64
      dataset = get_dataset(dataset="nations")
      r = TypedMessagePassingRepresentation(
          triples_factory=dataset.training,
          base_kwargs=dict(shape=embedding_dim),
          layers="rgcn",
          layers_kwargs=dict(
              in_channels=embedding_dim,
              out_channels=embedding_dim,
              num_bases=2,
              num_relations=dataset.num_relations,
          ),
      )
  
  ---
  name: Typed Message Passing
docstr: Return the (selected part of the) edge type.
  
  :param edge_mask: shape: `(num_edges,)`
      the edge mask
  
  :return: shape: `(num_selected_edges,)`
      the selected edge types, or all edge types if `edge_mask` is None
docstr_tok: ["Return","the","(","selected","part","of","the",")","edge","type",".",":","param","edge_mask",":","shape",":","`","(","num_edges",",",")","`","the","edge","mask",":","return",":","shape",":","`","(","num_selected_edges",",",")","`","the","selected","edge","types",",","or","all","edge","types","if","`","edge_mask","`","is","None"]
code: def _get_edge_type(
      self, edge_mask: Optional[torch.BoolTensor] = None
  ) -> torch.LongTensor:
      """
      Return the (selected part of the) edge type.
  
      :param edge_mask: shape: `(num_edges,)`
          the edge mask
  
      :return: shape: `(num_selected_edges,)`
          the selected edge types, or all edge types if `edge_mask` is None
      """
      if edge_mask is None:
          return self.edge_type
      return self.edge_type[edge_mask]
  
imports: from abc import ABC, abstractmethod
  from typing import Collection, Literal, Optional, Sequence
  import torch
  from class_resolver import ClassResolver, HintOrType, OneOrManyHintOrType, OneOrManyOptionalKwargs, OptionalKwargs
  from class_resolver.contrib.torch import activation_resolver
  from docdata import parse_docdata
  from torch import nn
  from .representation import Representation
  from .utils import ShapeError
  from ..triples.triples_factory import CoreTriplesFactory
  from ..typing import OneOrSequence
  from ..utils import get_edge_index, upgrade_to_sequence
label_desc: Use image node_id 2 for calling the TypedMessagePassingRepresentation obj's underlying member method code with example usage: obj._get_edge_type(edge_mask) and returns: self, self