neural_lam.models.hi_lam_parallel

neural_lam.models.hi_lam_parallel#

Parallel message-passing variant of the Hi-LAM architecture.

Module Contents#

class neural_lam.models.hi_lam_parallel.HiLAMParallel(args, config: neural_lam.config.NeuralLAMConfig, datastore: neural_lam.datastore.BaseDatastore)#

Bases: neural_lam.models.base_hi_graph_model.BaseHiGraphModel

Version of HiLAM where all message passing in the hierarchical mesh (up, down, inter-level) is ran in parallel.

This is a somewhat simpler alternative to the sequential message passing of Hi-LAM.

Initialize the parallel hierarchical message-passing processor.

hi_processor_step(mesh_rep_levels, mesh_same_rep, mesh_up_rep, mesh_down_rep)#

Internal processor step executed between mesh init and read-out.

Parameters:

  • mesh_rep_levels (list[torch.Tensor]) –

    Mesh representations for each level.

    • Shape: (B, N_mesh[l], d_h)

  • mesh_same_rep (list[torch.Tensor]) –

    Same-level edge representations.

    • Shape: (B, M_same[l], d_h)

  • mesh_up_rep (list[torch.Tensor]) –

    Upward edge representations.

    • Shape: (B, M_up[l], d_h)

  • mesh_down_rep (list[torch.Tensor]) –

    Downward edge representations.

    • Shape: (B, M_down[l], d_h)

Returns:

  • tuple[ – list[torch.Tensor], list[torch.Tensor], list[torch.Tensor], list[torch.Tensor]

  • ] – Updated representations (mesh_rep_levels, mesh_same_rep, mesh_up_rep, mesh_down_rep) after the parallel pass.

edge_split_sections#
Contents