deepmd_gnn.export
=================

.. py:module:: deepmd_gnn.export

.. autoapi-nested-parse::

   Helpers used by PyTorch-exportable model paths.



Functions
---------

.. autoapisummary::

   deepmd_gnn.export.pad_nlist_for_export
   deepmd_gnn.export.clear_export_guards_once


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

.. py:function:: pad_nlist_for_export(nlist: torch.Tensor) -> torch.Tensor

   Append a sentinel neighbor so symbolic export keeps neighbor axes dynamic.


.. py:function:: clear_export_guards_once(traced: torch.nn.Module) -> None

   Clear over-specialized guards from the next export of ``traced``.

   ``make_fx`` traces may specialize symbolic atom counts too aggressively.
   DeePMD's export path calls ``torch.export.export`` immediately after tracing,
   so this one-shot wrapper relaxes those constraints only for that export.