deepmd_gnn.edge
===============

.. py:module:: deepmd_gnn.edge

.. autoapi-nested-parse::

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



Functions
---------

.. autoapisummary::

   deepmd_gnn.edge._mm_tensor
   deepmd_gnn.edge._fake_dense_edge_index
   deepmd_gnn.edge.dense_edge_index


Module Contents
---------------

.. py:function:: _mm_tensor(mm_types: list[int]) -> torch.Tensor

.. py:function:: _fake_dense_edge_index(nlist: torch.Tensor, _atype: torch.Tensor, _mm: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]

.. py:function:: 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.