deepmd_gnn.edge

Neighbor-list C++ op wrappers used by graph model export paths.

Functions

_mm_tensor(→ torch.Tensor)
_fake_dense_edge_index(→ tuple[torch.Tensor, torch.Tensor])
dense_edge_index(→ tuple[torch.Tensor, torch.Tensor])	Build dense edge indices and validity mask with the C++ op.

Module Contents

deepmd_gnn.edge._mm_tensor(mm_types: list[int]) → torch.Tensor

deepmd_gnn.edge._fake_dense_edge_index(nlist: torch.Tensor, _atype: torch.Tensor, _mm: torch.Tensor) → tuple[torch.Tensor, torch.Tensor]

deepmd_gnn.edge.dense_edge_index(nlist: torch.Tensor, extended_atype: torch.Tensor, mm_types: list[int]) → tuple[torch.Tensor, torch.Tensor]

Build dense edge indices and validity mask with the C++ op.