deepmd.pt_expt.model.make_model

deepmd.pt_expt.model.make_model#

Classes#

_WrapperForwardEnergy

Callable wrapper for torch.autograd.functional.hessian.

Functions#

`_pad_nlist_for_export`(→ torch.Tensor)	Append a single `-1` column to `nlist` for export-time tracing.
`_cal_hessian_ext`(→ torch.Tensor)	Compute hessian of reduced output w.r.t. extended coordinates.
`make_model`() → type)	Make a model as a derived class of an atomic model.

Module Contents#

deepmd.pt_expt.model.make_model._pad_nlist_for_export(nlist: torch.Tensor) → torch.Tensor[source]#

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.

deepmd.pt_expt.model.make_model._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[source]#

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:

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

class deepmd.pt_expt.model.make_model._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)[source]#

Callable wrapper for torch.autograd.functional.hessian.

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

model[source]#

kk[source]#

ci[source]#

nall[source]#

atype[source]#

nlist[source]#

mapping[source]#

fparam[source]#

aparam[source]#

charge_spin = None[source]#

__call__(coord_flat: torch.Tensor) → torch.Tensor[source]#

Compute scalar reduced energy for one frame, one component.

Parameters:

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

Returns:

torch.Tensor: Scalar energy component.

deepmd.pt_expt.model.make_model.make_model(T_AtomicModel: type[deepmd.dpmodel.atomic_model.base_atomic_model.BaseAtomicModel], T_Bases: tuple[type, Ellipsis] = ()) → type[source]#

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:

CM: The model.