deepmd.jax.model.dp_model

Functions

make_jax_dp_model_from_dpmodel(...)

Make a JAX backend DP model from a DPModel backend DP model.

Module Contents

deepmd.jax.model.dp_model.make_jax_dp_model_from_dpmodel(dpmodel_model: type[deepmd.dpmodel.model.DPModelCommon], jax_atomicmodel: type[deepmd.jax.atomic_model.dp_atomic_model.DPAtomicModel]) → type[deepmd.dpmodel.model.DPModelCommon]

Make a JAX backend DP model from a DPModel backend DP model.

Parameters:

dpmodel_modeltype[DPModelCommon]

The DPModel backend DP model.

jax_atomicmodeltype[DPAtomicModel]

The JAX backend DP atomic model.

Returns:

type[DPModelCommon]

The JAX backend DP model.