pelican.nets.PELICAN

class pelican.nets.PELICAN(in_channels_rank2, in_channels_rank1, in_channels_rank0, out_rank, out_channels, num_blocks, hidden_channels, increase_hidden_channels=1.0, map_multipliers=True, factorize=False, activation='leaky_relu', dropout_prob=None, aggr='mean', compile=False, checkpoint_blocks=False)[source]

Bases: Module

PELICAN network.

PELICAN stands for Permutation Equivariant and Lorentz Invariant or Covariant Aggregator Network. PELICAN takes input of varying rank (0, 1, or 2), projects them onto rank 2 objects (edges), processes them with multiple PELICANBlock, and extracts output of varying rank (0, 1, or 2).

Parameters:

  • in_channels_rank2 (int) – Edge-level input features (rank 2). Can be zero.

  • in_channels_rank1 (int) – Node-level input features (rank 1). Can be zero.

  • in_channels_rank0 (int) – Graph-level input features (rank 0). Can be zero.

  • out_rank (int) – Rank of the output features (0, 1, or 2).

  • out_channels (int) – Number of output channels.

  • num_blocks (int) – Number of PELICAN blocks.

  • hidden_channels (int) – Number of hidden channels.

  • increase_hidden_channels (float) – Factor to increase hidden channels in the feedforward network. Default is 1.0.

  • map_multipliers (bool) – Whether to use learnable multipliers for each aggregation map. Default is True.

  • factorize (bool) – Whether to use factorized linear layers in the feedforward network. Default is False.

  • activation (str) – Activation function to use (‘gelu’, ‘relu’, ‘leaky_relu’, ‘tanh’, ‘sigmoid’, ‘silu’), by default ‘leaky_relu’.

  • dropout_prob (float) – Dropout probability in the feedforward network, by default None (no dropout).

  • aggr (str) – Aggregation method to use (‘mean’, ‘sum’, ‘prod’, ‘amin’, ‘amax’), by default ‘mean’.

  • compile (bool) – Whether to compile the model with torch.compile. Default is False. Compiling the model leads to significant speedups on GPU, because the aggregation functions involve many small operations that otherwise require many individual kernel launches. It is recommended to run model = torch.compile(model, **kwargs) outside of the constructor, however we provide this option for convenience. Note: When compile=True, the model requires the num_graphs argument in the forward pass to avoid a graph break.

  • checkpoint_blocks (bool) – Whether to use gradient checkpointing for PELICAN blocks to save memory. Default is False.

forward(edge_index, batch, in_rank2=None, in_rank1=None, in_rank0=None, num_graphs=None)[source]

Forward pass.

Parameters:

  • edge_index (torch.Tensor) – Edge index tensor of shape (2, E).

  • batch (torch.Tensor) – Batch tensor of shape (N).

  • in_rank2 (torch.Tensor) – Edge-level input features of shape (E, in_channels_rank2), by default None.

  • in_rank1 (torch.Tensor) – Node-level input features of shape (N, in_channels_rank1), by default None.

  • in_rank0 (torch.Tensor) – Graph-level input features of shape (G, in_channels_rank0), by default None.

  • num_graphs (int) – The number of graphs G in the batch, also known as batch size. If None, it will be inferred from the batch tensor. Inferring this number from the batch tensor requires a GPU/CPU synchronization, which slows down the code when running on GPU. Currently, the code requires the num_graphs argument in case compile=True.

Returns:

out – Output features of shape (G, out_channels) for out_rank=0, (N, out_channels) for out_rank=1, or (E, out_channels) for out_rank=2.

Return type:

torch.Tensor