deepmd.pt_expt.model.make_model
===============================

.. py:module:: deepmd.pt_expt.model.make_model


Classes
-------

.. autoapisummary::

   deepmd.pt_expt.model.make_model._WrapperForwardEnergy


Functions
---------

.. autoapisummary::

   deepmd.pt_expt.model.make_model._pad_nlist_for_export
   deepmd.pt_expt.model.make_model._cal_hessian_ext
   deepmd.pt_expt.model.make_model.make_model


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

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

   
   Append a single ``-1`` column to ``nlist`` for export-time tracing.

   Used inside ``forward_common_lower_exportable`` (and its spin counterpart)
   so that ``_format_nlist``'s terminal slice ``ret[..., :nnei]`` truncates
   to a statically sized output.  Without the extra column, torch.export
   cannot prove the ``ret.shape[-1] == nnei`` assertion at trace time and
   would specialise the dynamic ``nnei`` dim to the sample value.

   Combined with the short-circuit order in ``_format_nlist``
   (``extra_nlist_sort`` on the left) and the ``need_sorted_nlist_for_lower``
   override during tracing, this keeps the compiled graph's ``nnei`` axis
   fully dynamic and free of symbolic shape guards.















   ..
       !! processed by numpydoc !!

.. py:function:: _cal_hessian_ext(model: Any, kk: str, vdef: deepmd.dpmodel.output_def.OutputVariableDef, extended_coord: torch.Tensor, extended_atype: torch.Tensor, nlist: torch.Tensor, mapping: torch.Tensor | None, fparam: torch.Tensor | None, aparam: torch.Tensor | None, create_graph: bool = False, charge_spin: torch.Tensor | None = None) -> torch.Tensor

   
   Compute hessian of reduced output w.r.t. extended coordinates.

   Mirrors the JAX approach: compute hessian on extended coordinates,
   then let communicate_extended_output map nall->nloc.

   :Parameters:

       **model**
           The model (CM instance). Must have ``atomic_model.forward_common_atomic``.

       **kk**
           The output key (e.g. "energy").

       **vdef**
           The output variable definition.

       **extended_coord**
           Extended coordinates. Shape: [nf, nall, 3].

       **extended_atype**
           Extended atom types. Shape: [nf, nall].

       **nlist**
           Neighbor list. Shape: [nf, nloc, nsel].

       **mapping**
           Mapping from extended to local. Shape: [nf, nall] or None.

       **fparam**
           Frame parameters. Shape: [nf, nfp] or None.

       **aparam**
           Atomic parameters. Shape: [nf, nloc, nap] or None.

       **create_graph**
           Whether to create graph for higher-order derivatives.



   :Returns:

       :obj:`torch.Tensor`
           Hessian on extended coordinates. Shape: [nf, *def, nall, 3, nall, 3].











   ..
       !! processed by numpydoc !!

.. py:class:: _WrapperForwardEnergy(model: Any, kk: str, ci: int, nall: int, atype: torch.Tensor, nlist: torch.Tensor, mapping: torch.Tensor | None, fparam: torch.Tensor | None, aparam: torch.Tensor | None, charge_spin: torch.Tensor | None = None)

   
   Callable wrapper for torch.autograd.functional.hessian.

   Given flattened extended coordinates, recomputes the reduced energy
   for one frame and one output component.















   ..
       !! processed by numpydoc !!

   .. py:attribute:: model


   .. py:attribute:: kk


   .. py:attribute:: ci


   .. py:attribute:: nall


   .. py:attribute:: atype


   .. py:attribute:: nlist


   .. py:attribute:: mapping


   .. py:attribute:: fparam


   .. py:attribute:: aparam


   .. py:attribute:: charge_spin
      :value: None



   .. py:method:: __call__(coord_flat: torch.Tensor) -> torch.Tensor

      
      Compute scalar reduced energy for one frame, one component.


      :Parameters:

          **coord_flat**
              Flattened extended coordinates for one frame. Shape: [nall * 3].



      :Returns:

          :obj:`torch.Tensor`
              Scalar energy component.











      ..
          !! processed by numpydoc !!


.. py:function:: make_model(T_AtomicModel: type[deepmd.dpmodel.atomic_model.base_atomic_model.BaseAtomicModel], T_Bases: tuple[type, Ellipsis] = ()) -> type

   
   Make a model as a derived class of an atomic model.

   Wraps dpmodel's make_model with torch.nn.Module and overrides
   forward_common_atomic to use autograd-based derivatives.

   :Parameters:

       **T_AtomicModel**
           The atomic model.

       **T_Bases**
           Additional base classes for the returned model class.
           For example, pass ``(BaseModel,)`` so that the concrete model
           inherits the pt_expt ``BaseModel`` plugin registry.



   :Returns:

       :obj:`CM`
           The model.











   ..
       !! processed by numpydoc !!