model

Models

Chemtrain does not require a specific model to be used. Nevertheless, it provides some pre-defined models, such as a simple classical force field in the jax_md_mod.model.prior module.

• model.prior
• model.neural_networks

Building Blocks

These submodules contain utilities necessary for building advanced potential models. For example, the jax_md_mod.model.sparse_graph module contains functions to construct a sparse graph from a JAX, M.D jax_md.partition.NeighborList.

• model.dimnet_basis_util
• model.layers
• model.sparse_graph