deepmd 0.1.0

DeePMD-kit deep potential models as RLX IR graph builders
Documentation 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106

# deepmd

Repository: [github.com/eugenehp/deepmd-rs](https://github.com/eugenehp/deepmd-rs)

DeePMD-kit deep potential models as [RLX](https://docs.rs/rlx) IR graph builders.

This crate translates the inference-side graph construction from
[`deepmd.dpmodel`](https://github.com/deepmodeling/deepmd-kit) into
[`rlx_ir::Graph`](https://docs.rs/rlx-ir) objects. Compile and execute
them on any RLX backend (CPU, Metal, MLX, CUDA, WGPU, …).

## Quick start

Add to `Cargo.toml`:

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

Build and run a small energy model on CPU:

```rust
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`):

| Feature | Backend | Platform |
|---------|---------|----------|
| *(default)* | CPU | everywhere |
| `metal` | Apple Metal | macOS |
| `mlx` | Apple MLX | macOS |
| `cuda` | NVIDIA CUDA | Linux / Windows + NVIDIA GPU |
| `wgpu` | cross-vendor GPU | Metal / Vulkan / DX12 |
| `ane` | Apple Neural Engine | macOS (requires CoreML backend in RLX) |

```bash
cargo build -p deepmd --features metal,mlx
```

See [docs/backend-selection.md](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:

```bash
# 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.