deepmd

Repository: github.com/eugenehp/deepmd-rs

DeePMD-kit deep potential models as RLX IR graph builders.

This crate translates the inference-side graph construction from deepmd.dpmodel into rlx_ir::Graph objects. Compile and execute them on any RLX backend (CPU, Metal, MLX, CUDA, WGPU, …).

Quick start

Add to Cargo.toml:

[dependencies]
deepmd = "0.1"
rlx-runtime = { version = "0.2", default-features = false, features = ["cpu"] }

Build and run a small energy model on CPU:

use deepmd::config::{DPModelConfig, EnerFittingConfig, SeAConfig};
use deepmd::model::build_dp_energy_graph;
use rlx_ir::DType;
use rlx_runtime::{Device, Session};

let cfg = DPModelConfig {
    descriptor: SeAConfig {
        rcut: 6.0,
        rcut_smth: 0.5,
        sel: vec![4],
        neuron: vec![8, 16, 32],
        axis_neuron: 4,
        resnet_dt: false,
        type_one_side: true,
        activation_function: "tanh".into(),
        type_map: None,
    },
    fitting_net: EnerFittingConfig {
        ntypes: 1,
        dim_descrpt: 32,
        dim_out: 1,
        neuron: vec![120, 120, 120],
        resnet_dt: false,
        numb_fparam: 0,
        numb_aparam: 0,
        dim_case_embd: 0,
        activation_function: "tanh".into(),
        mixed_types: true,
    },
    type_map: None,
};

let nf = 1;
let nloc = 4;
let (graph, handle) = build_dp_energy_graph(&cfg, nf, nloc).unwrap();
let mut graph = graph;
graph.set_outputs(vec![handle.energy]);

let session = Session::new(Device::Cpu);
let compiled = session.compile(graph);
// Bind weights with `compiled.set_param_typed(...)`, then `compiled.run(...)`.
let _ = compiled;

Optional backends

Enable GPU backends via Cargo features on deepmd (they forward to rlx-runtime):

cargo build -p deepmd --features metal,mlx

See docs/backend-selection.md for measured latency crossovers and parity notes across backends.

Testing

Golden parity tests compare RLX execution against Python deepmd-kit reference fixtures:

# CPU parity (default)
cargo test --release -p deepmd

# Cross-backend sweep (enable features for your hardware)
cargo test --release -p deepmd --features metal,mlx,wgpu \
    --test parity_cross_backend -- --nocapture --test-threads=1

License

GPL-3.0-only. See the RLX project for upstream licensing details.