def METHOD_NAME(
    self,
    x: Tensor,
    index: Optional[Tensor] = None,
    ptr: Optional[Tensor] = None,
    dim_size: Optional[int] = None,
    dim: int = -2,
    max_num_elements: Optional[int] = None,
) -> Tensor:
    x, mask = self.to_dense_batch(x, index, ptr, dim_size, dim,
                                  max_num_elements=max_num_elements)
    for encoder in self.encoders:
        x = encoder(x, mask)
    x = self.pma(x, mask)
    for decoder in self.decoders:
        x = decoder(x)
    x = x.nan_to_num()
    return x.flatten(1, 2) if self.concat else x.mean(dim=1)