deepmd.jax.model.base_model

Attributes

BaseModel

Functions

forward_common_atomic(self, extended_coord, ...[, ...])

Module Contents

deepmd.jax.model.base_model.BaseModel

deepmd.jax.model.base_model.forward_common_atomic(self, extended_coord: deepmd.jax.env.jnp.ndarray, extended_atype: deepmd.jax.env.jnp.ndarray, nlist: deepmd.jax.env.jnp.ndarray, mapping: deepmd.jax.env.jnp.ndarray | None = None, fparam: deepmd.jax.env.jnp.ndarray | None = None, aparam: deepmd.jax.env.jnp.ndarray | None = None, do_atomic_virial: bool = False)