deepmd 0.1.0

// RLX — versatile ML compiler + runtime.
// Copyright (C) 2026 Eugene Hauptmann, Nataliya Kosmyna.

//! Repformer descriptor block — DPA-2's message-passing layer stack.
//!
//! Translated from `deepmd/dpmodel/descriptor/repformers.py`.
//!
//! ## Architecture
//!
//! Three per-atom and per-pair feature channels are evolved through
//! `nlayers` `RepformerLayer`s:
//!
//! * `g1 ∈ ℝ^{nf, nloc, ng1}`     — per-atom invariant
//! * `g2 ∈ ℝ^{nf, nloc, nnei, ng2}` — per-pair invariant
//! * `h2 ∈ ℝ^{nf, nloc, nnei, 3}`  — per-pair equivariant
//!
//! Each layer optionally runs:
//!
//! * `Atten2Map(g2, h2) → AAg`            — pair-pair attention map
//! * `Atten2MultiHeadApply(AAg, g2) → g2'` — apply to g2
//! * `LocalAtten(g1, gg1) → g1'`           — local attention on neighbors
//! * `symmetrization_op(g2/gg1, h2) → grrg/drrd` — invariant axial projection
//! * MLPs on g1 / g2 / h2
//!
//! ## Parity status
//!
//! This port targets the **mainstream config** of
//! `DescrptBlockRepformers` (`smooth=true`, `update_style="res_avg"`,
//! `use_sqrt_nnei=true`, `g1_out_conv=true`, `g1_out_mlp=true`, default
//! widths).  Non-default branches return an error from
//! [`build_repformers`].

use anyhow::{bail, Result};
use rlx_ir::infer::GraphExt;
use rlx_ir::op::{Activation, BinaryOp, ReduceOp};
use rlx_ir::{DType, Graph, NodeId, Shape};
use serde::{Deserialize, Serialize};

use crate::nn::{scalar_const, ActivationKind};

/// Top-level repformer-block config (RepformerArgs in Python).
#[derive(Debug, Clone, Deserialize, Serialize)]
pub struct RepformersConfig {
    pub rcut: f64,
    pub rcut_smth: f64,
    pub sel: usize,
    pub ntypes: usize,
    #[serde(default = "default_nlayers")]
    pub nlayers: usize,
    #[serde(default = "default_g1_dim")]
    pub g1_dim: usize,
    #[serde(default = "default_g2_dim")]
    pub g2_dim: usize,
    #[serde(default = "default_axis_neuron")]
    pub axis_neuron: usize,
    #[serde(default = "default_true")]
    pub update_g1_has_conv: bool,
    #[serde(default = "default_true")]
    pub update_g1_has_drrd: bool,
    #[serde(default = "default_true")]
    pub update_g1_has_grrg: bool,
    #[serde(default = "default_true")]
    pub update_g1_has_attn: bool,
    #[serde(default = "default_true")]
    pub update_g2_has_g1g1: bool,
    #[serde(default = "default_true")]
    pub update_g2_has_attn: bool,
    #[serde(default)]
    pub update_h2: bool,
    #[serde(default = "default_attn1_hidden")]
    pub attn1_hidden: usize,
    #[serde(default = "default_attn1_nhead")]
    pub attn1_nhead: usize,
    #[serde(default = "default_attn2_hidden")]
    pub attn2_hidden: usize,
    #[serde(default = "default_attn2_nhead")]
    pub attn2_nhead: usize,
    #[serde(default)]
    pub attn2_has_gate: bool,
    #[serde(default = "default_activation")]
    pub activation_function: String,
    #[serde(default = "default_update_style")]
    pub update_style: String,
    #[serde(default = "default_true")]
    pub smooth: bool,
    #[serde(default = "default_true")]
    pub use_sqrt_nnei: bool,
    #[serde(default = "default_true")]
    pub g1_out_conv: bool,
    #[serde(default = "default_true")]
    pub g1_out_mlp: bool,
    #[serde(default = "default_ln_eps")]
    pub ln_eps: f32,
    #[serde(default = "default_epsilon")]
    pub epsilon: f32,
}

fn default_nlayers() -> usize {
    3
}
fn default_g1_dim() -> usize {
    128
}
fn default_g2_dim() -> usize {
    16
}
fn default_axis_neuron() -> usize {
    4
}
fn default_attn1_hidden() -> usize {
    64
}
fn default_attn1_nhead() -> usize {
    4
}
fn default_attn2_hidden() -> usize {
    16
}
fn default_attn2_nhead() -> usize {
    4
}
fn default_activation() -> String {
    "tanh".into()
}
fn default_update_style() -> String {
    "res_avg".into()
}
fn default_true() -> bool {
    true
}
fn default_ln_eps() -> f32 {
    1e-5
}
fn default_epsilon() -> f32 {
    1e-4
}

impl RepformersConfig {
    pub fn ng1(&self) -> usize {
        self.g1_dim
    }
    pub fn ng2(&self) -> usize {
        self.g2_dim
    }
}

pub struct RepformersInputs {
    /// `[nf, nall, ng1]` — initial per-atom (extended) embedding, fed
    /// by the repinit block.
    pub g1_ext: NodeId,
    /// `[nf, nloc, nnei, 4]` — geometric env matrix from host.
    pub env_mat_raw: NodeId,
    /// `[nf, nloc, nnei]` i32.
    pub nlist: NodeId,
    /// `[nf, nloc, nnei]` f32 mask.
    pub nlist_mask: NodeId,
    /// `[nf, nloc, nnei]` switching weight.
    pub sw: NodeId,
    /// `[nf, nloc]` i32 local atom types.
    pub atype_loc: NodeId,
    /// `[nf, nall]` i32 mapping that resolves each extended atom to its
    /// owning local atom.  Used between repformer layers to rebuild
    /// `g1_ext` from the updated local `g1`.  Required when `nall > nloc`;
    /// optional when `nall == nloc` (the identity case).
    pub mapping: Option<NodeId>,
}

pub struct RepformersOutputs {
    pub g1: NodeId, // [nf, nloc, ng1]
    pub g2: NodeId, // [nf, nloc, nnei, ng2]
    pub h2: NodeId, // [nf, nloc, nnei, 3]
    pub rot_mat: NodeId, // [nf, nloc, ng2, 3]
    pub sw: NodeId,
}

pub fn build_repformers(
    g: &mut Graph,
    cfg: &RepformersConfig,
    inputs: RepformersInputs,
    nf: usize,
    nloc: usize,
    nall: usize,
) -> Result<RepformersOutputs> {
    if cfg.update_style != "res_avg" {
        bail!(
            "repformers: only update_style=\"res_avg\" is implemented (got \"{}\")",
            cfg.update_style
        );
    }
    let activation = ActivationKind::parse(&cfg.activation_function)?;
    let _ng1 = cfg.ng1();
    let ng2 = cfg.ng2();
    let nnei = cfg.sel;

    // ── davg/dstd normalization on env mat ──
    let davg = g.param(
        "repformers.davg",
        Shape::new(&[cfg.ntypes, nnei, 4], DType::F32),
    );
    let dstd = g.param(
        "repformers.dstd",
        Shape::new(&[cfg.ntypes, nnei, 4], DType::F32),
    );
    let davg_g = g.gather_(davg, inputs.atype_loc, 0);
    let dstd_g = g.gather_(dstd, inputs.atype_loc, 0);
    let mut dmatrix = g.sub(inputs.env_mat_raw, davg_g);
    dmatrix = g.div(dmatrix, dstd_g);

    // g2_init = act(linear(s(r)))  where s(r) = dmatrix[..., 0:1].
    let radial = g.narrow_(dmatrix, 3, 0, 1); // [nf, nloc, nnei, 1]
    let g2_w = g.param(
        "repformers.g2_embd.w",
        Shape::new(&[1, ng2], DType::F32),
    );
    let g2_b = g.param("repformers.g2_embd.b", Shape::new(&[ng2], DType::F32));
    let g2_pre = g.mm(radial, g2_w);
    let g2_pre = g.add(g2_pre, g2_b);
    let mut g2 = crate::nn::apply_activation(g, activation, g2_pre);

    // h2_init = dmatrix[..., 1:4]
    let mut h2 = g.narrow_(dmatrix, 3, 1, 3); // [nf, nloc, nnei, 3]

    // g1_init = act(atype_embd) — Python applies act to the whole extended
    // embedding before per-layer ghost exchange, so gg1 = gather(g1_ext, nlist)
    // inside the layer sees activated values.  For nall==nloc identity mapping
    // this is equivalent to activating g1_ext pointwise.
    let g1_ext_act = crate::nn::apply_activation(g, activation, inputs.g1_ext);
    let mut g1 = g.narrow_(g1_ext_act, 1, 0, nloc); // [nf, nloc, ng1]

    // Apply nlist_mask · sw on g2.
    let mask_4d_shape = Shape::new(&[nf, nloc, nnei, 1], DType::F32);
    let mask_4d = g.reshape(
        inputs.nlist_mask,
        vec![nf as i64, nloc as i64, nnei as i64, 1],
        mask_4d_shape,
    );
    g2 = g.mul(g2, mask_4d);

    // Per-layer Python recomputes `g1_ext = _exchange_ghosts(g1)` so
    // `gg1 = gather(g1_ext, nlist)` sees the just-updated local g1.
    // For nall==nloc identity mapping the exchange is a no-op (g1_ext = g1).
    // For nall>nloc we rebuild via mapping with `gather(g1, mapping, axis=1)`.
    let mut g1_ext_cur = g1_ext_act;

    for l in 0..cfg.nlayers {
        // Python sets `update_chnnl_2 = (l != nlayers - 1)`; the last
        // layer only updates g1.
        let is_last = l == cfg.nlayers - 1;
        let prefix = format!("repformers.layer{l}");
        let (g1_new, g2_new, h2_new) = build_repformer_layer(
            g, cfg, &prefix, activation, g1, g2, h2, inputs.nlist, inputs.nlist_mask,
            inputs.sw, g1_ext_cur, nf, nloc, nnei, !is_last,
        )?;
        g1 = g1_new;
        g2 = g2_new;
        h2 = h2_new;
        // Rebuild g1_ext for the next iteration from the updated g1.
        if l + 1 < cfg.nlayers {
            g1_ext_cur = if nall == nloc {
                g1
            } else {
                let mapping = inputs.mapping.ok_or_else(|| {
                    anyhow::anyhow!("repformers: mapping is required when nall != nloc")
                })?;
                g.gather_(g1, mapping, 1)
            };
        }
    }

    // rot_mat = hg^T where hg = (1/√nnei) · h^T · g2 from _cal_hg.
    let hg = build_cal_hg(g, g2, h2, inputs.sw, cfg, nf, nloc, nnei)?; // [nf, nloc, 3, ng2]
    let rot_mat = g.transpose_(hg, vec![0, 1, 3, 2]); // [nf, nloc, ng2, 3]

    Ok(RepformersOutputs {
        g1,
        g2,
        h2,
        rot_mat,
        sw: inputs.sw,
    })
}

/// One repformer layer with the mainstream update path:
/// `g1' = act(linear1(concat([g1, conv, grrg, drrd, attn]))) + g1`
/// `g2' = act(linear2(g2)) + g2 + g1g1_proj + attn`
/// `h2' = h2 + linear_head(AAg · h2)` (if `update_h2`)
fn build_repformer_layer(
    g: &mut Graph,
    cfg: &RepformersConfig,
    prefix: &str,
    activation: ActivationKind,
    g1: NodeId,
    g2: NodeId,
    h2: NodeId,
    nlist: NodeId,
    nlist_mask: NodeId,
    sw: NodeId,
    g1_ext: NodeId,
    nf: usize,
    nloc: usize,
    nnei: usize,
    update_chnnl_2: bool,
) -> Result<(NodeId, NodeId, NodeId)> {
    let ng1 = cfg.g1_dim;
    let ng2 = cfg.g2_dim;
    let axis = cfg.axis_neuron;

    // gg1 = neighbor g1 — gather g1_ext by nlist along axis 1.
    let gg1 = make_nei_g1(g, g1_ext, nlist, nf, nloc, nnei, ng1);

    // Python disables g2 / h2 updates on the last layer
    // (`update_chnnl_2 = False`).  Mirror that here.
    let do_chnnl_2 = update_chnnl_2;
    let do_g2_g1g1 = do_chnnl_2 && cfg.update_g2_has_g1g1;
    let do_g2_attn = do_chnnl_2 && cfg.update_g2_has_attn;
    let do_h2 = do_chnnl_2 && cfg.update_h2;

    // ── g2 update ──
    let (g2_new, _aag_dbg) = if do_chnnl_2 {
        let g2_1 = linear_act(g, &format!("{prefix}.linear2"), g2, ng2, ng2, activation);
        let g2_g1g1 = if do_g2_g1g1 {
            let g1g1 = update_g2_g1g1(g, g1, gg1, nlist_mask, sw, nf, nloc, nnei, ng1);
            Some(linear(g, &format!("{prefix}.proj_g1g1g2"), g1g1, ng1, ng2))
        } else {
            None
        };
        let aag = if do_g2_attn || do_h2 {
            Some(build_atten2_map(
                g, cfg, &format!("{prefix}.attn2g_map"), g2, h2, nlist_mask, sw,
                nf, nloc, nnei,
            )?)
        } else {
            None
        };
        let g2_attn = if do_g2_attn {
            let g2_mh = build_atten2_mh_apply(
                g, cfg, &format!("{prefix}.attn2_mh_apply"), aag.unwrap(), g2,
                nf, nloc, nnei, ng2,
            );
            // Python: `g2_2 = self.attn2_lm(self.attn2_mh_apply(...))` — LN over the
            // last (ng2) axis with learnable gamma/beta.
            let gamma = g.param(
                format!("{prefix}.attn2_lm.w"),
                Shape::new(&[ng2], DType::F32),
            );
            let beta = g.param(
                format!("{prefix}.attn2_lm.b"),
                Shape::new(&[ng2], DType::F32),
            );
            Some(g.ln(g2_mh, gamma, beta, cfg.ln_eps))
        } else {
            None
        };
        let mut g2_terms = vec![g2, g2_1];
        if let Some(t) = g2_g1g1 {
            g2_terms.push(t);
        }
        if let Some(t) = g2_attn {
            g2_terms.push(t);
        }
        (list_update_res_avg(g, &g2_terms), aag)
    } else {
        (g2, None)
    };

    // ── h2 update ──
    let h2_new = if do_h2 {
        let aag = _aag_dbg.expect("aag must be set if update_h2");
        let h2_upd = update_h2(g, &format!("{prefix}.attn2_ev_apply"), h2, aag, nf, nloc, nnei);
        list_update_res_avg(g, &[h2, h2_upd])
    } else {
        h2
    };

    // ── g1 update ──
    // Build the input concat for the main mlp:
    //   if g1_out_mlp:    only [conv, grrg, drrd] (g1 itself joins update list directly)
    //   else:             [g1, conv, grrg, drrd]
    let mut g1_mlp: Vec<NodeId> = Vec::new();
    if !cfg.g1_out_mlp {
        g1_mlp.push(g1);
    }

    let g1_conv = if cfg.update_g1_has_conv {
        Some(update_g1_conv(
            g, &format!("{prefix}"), cfg, gg1, g2, nlist_mask, sw, nf, nloc, nnei, ng1, ng2,
        ))
    } else {
        None
    };
    if !cfg.g1_out_conv {
        if let Some(c) = g1_conv {
            g1_mlp.push(c);
        }
    }

    if cfg.update_g1_has_grrg {
        let grrg =
            symmetrization_op(g, g2, h2, nlist_mask, sw, cfg, axis, nf, nloc, nnei, ng2);
        g1_mlp.push(grrg);
    }
    if cfg.update_g1_has_drrd {
        let drrd =
            symmetrization_op(g, gg1, h2, nlist_mask, sw, cfg, axis, nf, nloc, nnei, ng1);
        g1_mlp.push(drrd);
    }

    // concat → MLP linear1 → act → [nf, nloc, ng1]
    let concat_dim: usize = g1_mlp
        .iter()
        .map(|n| match g.shape(*n).dim(2) {
            rlx_ir::Dim::Static(d) => d,
            _ => 0,
        })
        .sum();
    let mlp_in_shape = Shape::new(&[nf, nloc, concat_dim], DType::F32);
    let g1_mlp_concat = g.concat(g1_mlp, 2, mlp_in_shape);
    let g1_1 = linear_act(
        g,
        &format!("{prefix}.linear1"),
        g1_mlp_concat,
        concat_dim,
        ng1,
        activation,
    );

    // Match Python's exact append order: [g1, g1_self_mlp?, conv?, g1_1, loc_attn?].
    // f32 addition isn't associative, so summation order leaks into the
    // last few bits.  Doing it in the same order as upstream eliminates
    // residual drift in the repformer-layer parity tests.
    let mut g1_updates = vec![g1];
    if cfg.g1_out_mlp {
        let g1_self = linear_act(
            g,
            &format!("{prefix}.g1_self_mlp"),
            g1,
            ng1,
            ng1,
            activation,
        );
        g1_updates.push(g1_self);
    }
    if cfg.g1_out_conv {
        if let Some(c) = g1_conv {
            g1_updates.push(c);
        }
    }
    g1_updates.push(g1_1);
    if cfg.update_g1_has_attn {
        let loc_attn = build_local_atten(
            g, cfg, &format!("{prefix}.loc_attn"), g1, gg1, nlist_mask, sw,
            nf, nloc, nnei, ng1,
        )?;
        g1_updates.push(loc_attn);
    }
    let g1_new = list_update_res_avg(g, &g1_updates);

    Ok((g1_new, g2_new, h2_new))
}

// ── helpers ──────────────────────────────────────────────────────────

fn linear(
    g: &mut Graph,
    prefix: &str,
    x: NodeId,
    in_dim: usize,
    out_dim: usize,
) -> NodeId {
    let w = g.param(
        format!("{prefix}.w"),
        Shape::new(&[in_dim, out_dim], DType::F32),
    );
    let b = g.param(format!("{prefix}.b"), Shape::new(&[out_dim], DType::F32));
    let mm = g.mm(x, w);
    g.add(mm, b)
}

fn linear_act(
    g: &mut Graph,
    prefix: &str,
    x: NodeId,
    in_dim: usize,
    out_dim: usize,
    activation: ActivationKind,
) -> NodeId {
    let pre = linear(g, prefix, x, in_dim, out_dim);
    crate::nn::apply_activation(g, activation, pre)
}

fn make_nei_g1(
    g: &mut Graph,
    g1_ext: NodeId,
    nlist: NodeId,
    nf: usize,
    nloc: usize,
    nnei: usize,
    ng1: usize,
) -> NodeId {
    // Mask negative indices to 0 (host should already do that, but be safe).
    // Flatten nlist to [nf, nloc*nnei], then gather over axis 1.
    let nlist_2d_shape = Shape::new(&[nf, nloc * nnei], DType::I32);
    let nlist_2d = g.reshape(
        nlist,
        vec![nf as i64, (nloc * nnei) as i64],
        nlist_2d_shape,
    );
    // gather(g1_ext, nlist_2d, axis=1) → [nf, nloc*nnei, ng1]
    // gather_shape: indices_dims + table.dims[axis+1..]
    //  = [nf, nloc*nnei] + [ng1] = [nf, nloc*nnei, ng1]
    let gathered = g.gather_(g1_ext, nlist_2d, 1);
    let out_shape = Shape::new(&[nf, nloc, nnei, ng1], DType::F32);
    g.reshape(
        gathered,
        vec![nf as i64, nloc as i64, nnei as i64, ng1 as i64],
        out_shape,
    )
}

fn update_g2_g1g1(
    g: &mut Graph,
    g1: NodeId, // [nf, nloc, ng1]
    gg1: NodeId, // [nf, nloc, nnei, ng1]
    nlist_mask: NodeId,
    sw: NodeId,
    nf: usize,
    nloc: usize,
    nnei: usize,
    ng1: usize,
) -> NodeId {
    // g1 broadcast to [nf, nloc, 1, ng1] then mul with gg1
    let g1_4d_shape = Shape::new(&[nf, nloc, 1, ng1], DType::F32);
    let g1_4d = g.reshape(
        g1,
        vec![nf as i64, nloc as i64, 1, ng1 as i64],
        g1_4d_shape,
    );
    let prod = g.mul(g1_4d, gg1); // [nf, nloc, nnei, ng1]
    // mask + sw multiply
    let mask_4d_shape = Shape::new(&[nf, nloc, nnei, 1], DType::F32);
    let mask_4d = g.reshape(
        nlist_mask,
        vec![nf as i64, nloc as i64, nnei as i64, 1],
        mask_4d_shape,
    );
    let masked = g.mul(prod, mask_4d);
    let sw_4d_shape = Shape::new(&[nf, nloc, nnei, 1], DType::F32);
    let sw_4d = g.reshape(
        sw,
        vec![nf as i64, nloc as i64, nnei as i64, 1],
        sw_4d_shape,
    );
    g.mul(masked, sw_4d)
}

fn update_g1_conv(
    g: &mut Graph,
    prefix: &str,
    cfg: &RepformersConfig,
    gg1: NodeId, // [nf, nloc, nnei, ng1]
    g2: NodeId,  // [nf, nloc, nnei, ng2]
    nlist_mask: NodeId,
    sw: NodeId,
    nf: usize,
    nloc: usize,
    nnei: usize,
    ng1: usize,
    ng2: usize,
) -> NodeId {
    // Mirror Python's `_update_g1_conv` (g1_out_conv=True branch):
    //   gg1 = reshape(gg1, [nf, nloc, nnei, ng1])     # already that shape
    //   gg1 = _apply_nlist_mask(gg1, nlist_mask)
    //   gg1 = _apply_switch(gg1, sw)                  # smooth=True
    //   invnnei = (1/nnei) * ones((nf, nloc, 1))
    //   g2_proj = proj_g1g2(g2)                       # [nf, nloc, nnei, ng1], bias=False
    //   g1_11 = sum(g2_proj * gg1, axis=2) * invnnei
    let mask_4d = g.reshape(
        nlist_mask,
        vec![nf as i64, nloc as i64, nnei as i64, 1],
        Shape::new(&[nf, nloc, nnei, 1], DType::F32),
    );
    let sw_4d = g.reshape(
        sw,
        vec![nf as i64, nloc as i64, nnei as i64, 1],
        Shape::new(&[nf, nloc, nnei, 1], DType::F32),
    );
    let gg1_m = g.mul(gg1, mask_4d);
    let gg1_ms = if cfg.smooth { g.mul(gg1_m, sw_4d) } else { gg1_m };

    // proj_g1g2: NativeLayer(g2_dim, g1_dim, bias=False)
    let proj_w = g.param(
        format!("{prefix}.proj_g1g2.w"),
        Shape::new(&[ng2, ng1], DType::F32),
    );
    let g2_proj = g.mm(g2, proj_w); // [nf, nloc, nnei, ng1]

    let prod = g.mul(g2_proj, gg1_ms); // [nf, nloc, nnei, ng1]
    let sum_shape = Shape::new(&[nf, nloc, ng1], DType::F32);
    let summed = g.reduce(prod, ReduceOp::Sum, vec![2], false, sum_shape);
    let inv_n = scalar_const(g, 1.0 / nnei as f32);
    g.mul(summed, inv_n)
}

fn update_h2(
    g: &mut Graph,
    prefix: &str,
    h2: NodeId,
    aag: NodeId, // [nf, nloc, nnei, nnei, nh]
    nf: usize,
    nloc: usize,
    nnei: usize,
) -> NodeId {
    // Mirror `Atten2EquiVarApply.call`:
    //   AA  = transpose_01423(AA)               # [nf, nloc, nh, nnei, nnei]
    //   h2m = broadcast(h2, axis=2 → nh)        # [nf, nloc, nh, nnei, 3]
    //   ret = matmul(AA, h2m)                   # [nf, nloc, nh, nnei, 3]
    //   ret = transpose_01342(ret)              # [nf, nloc, nnei, 3, nh]
    //   ret = head_map(ret)                     # NativeLayer(nh, 1, bias=False) → [..., 1]
    //   return squeeze(ret, -1)                 # [nf, nloc, nnei, 3]
    let nh = match g.shape(aag).dim(4) {
        rlx_ir::Dim::Static(n) => n,
        _ => panic!("aag head dim must be static"),
    };
    let aag_perm = g.transpose_(aag, vec![0, 1, 4, 2, 3]); // [nf, nloc, nh, nnei, nnei]
    // Tile h2 across nh: reshape to [nf, nloc, 1, nnei, 3] then concat would be heavy;
    // use mm directly — broadcasting in the 5D matmul: lhs [nf, nloc, nh, nnei, nnei],
    // rhs [nf, nloc, 1, nnei, 3] should broadcast over nh. Use reshape + repeat-via-mul.
    let h2_5d_shape = Shape::new(&[nf, nloc, 1, nnei, 3], DType::F32);
    let h2_5d = g.reshape(
        h2,
        vec![nf as i64, nloc as i64, 1, nnei as i64, 3],
        h2_5d_shape,
    );
    // Broadcast tile via concat (clones along nh axis) — cheap for small nh.
    let mut h2_tiles = Vec::with_capacity(nh);
    for _ in 0..nh {
        h2_tiles.push(h2_5d);
    }
    let h2_tiled_shape = Shape::new(&[nf, nloc, nh, nnei, 3], DType::F32);
    let h2_tiled = g.concat(h2_tiles, 2, h2_tiled_shape);
    let ret = g.mm(aag_perm, h2_tiled); // [nf, nloc, nh, nnei, 3]
    // transpose_01342: [0,1,3,4,2]
    let ret_perm = g.transpose_(ret, vec![0, 1, 3, 4, 2]); // [nf, nloc, nnei, 3, nh]
    // head_map: NativeLayer(nh, 1, bias=False) — applied to last dim
    let head_w = g.param(
        format!("{prefix}.head_map.w"),
        Shape::new(&[nh, 1], DType::F32),
    );
    let mapped = g.mm(ret_perm, head_w); // [nf, nloc, nnei, 3, 1]
    let out_shape = Shape::new(&[nf, nloc, nnei, 3], DType::F32);
    g.reshape(
        mapped,
        vec![nf as i64, nloc as i64, nnei as i64, 3],
        out_shape,
    )
}

fn build_cal_hg(
    g: &mut Graph,
    val: NodeId, // [nf, nloc, nnei, ng] (invariant)
    h: NodeId,   // [nf, nloc, nnei, 3]
    sw: NodeId,
    cfg: &RepformersConfig,
    nf: usize,
    nloc: usize,
    nnei: usize,
) -> Result<NodeId> {
    let _ = nnei;
    // mask + sw
    let sw_4d = g.reshape(
        sw,
        vec![nf as i64, nloc as i64, nnei as i64, 1],
        Shape::new(&[nf, nloc, nnei, 1], DType::F32),
    );
    let g_masked = if cfg.smooth { g.mul(val, sw_4d) } else { val };

    // hg = (h^T · g_masked) / √nnei
    let h_t = g.transpose_(h, vec![0, 1, 3, 2]); // [nf, nloc, 3, nnei]
    let hg = g.mm(h_t, g_masked); // [nf, nloc, 3, ng]
    let inv = if cfg.use_sqrt_nnei {
        scalar_const(g, 1.0 / (nnei as f32).sqrt())
    } else {
        scalar_const(g, 1.0 / nnei as f32)
    };
    Ok(g.mul(hg, inv))
}

fn symmetrization_op(
    g: &mut Graph,
    val: NodeId, // [nf, nloc, nnei, ng]
    h: NodeId,   // [nf, nloc, nnei, 3]
    nlist_mask: NodeId,
    sw: NodeId,
    cfg: &RepformersConfig,
    axis: usize,
    nf: usize,
    nloc: usize,
    nnei: usize,
    ng: usize,
) -> NodeId {
    let _ = nlist_mask;
    let hg = build_cal_hg(g, val, h, sw, cfg, nf, nloc, nnei).unwrap(); // [nf, nloc, 3, ng]
    // grrg = (hg[..., :axis]^T · hg) / 3 → [nf, nloc, axis, ng]
    let hgm = g.narrow_(hg, 3, 0, axis); // [nf, nloc, 3, axis]
    let hgm_t = g.transpose_(hgm, vec![0, 1, 3, 2]); // [nf, nloc, axis, 3]
    let grrg = g.mm(hgm_t, hg); // [nf, nloc, axis, ng]
    let inv_3 = scalar_const(g, 1.0 / 3.0);
    let grrg = g.mul(grrg, inv_3);
    let out_shape = Shape::new(&[nf, nloc, axis * ng], DType::F32);
    g.reshape(grrg, vec![nf as i64, nloc as i64, (axis * ng) as i64], out_shape)
}

fn list_update_res_avg(g: &mut Graph, list: &[NodeId]) -> NodeId {
    let n = list.len();
    let mut acc = list[0];
    for &x in &list[1..] {
        acc = g.add(acc, x);
    }
    let inv_sqrt = scalar_const(g, 1.0 / (n as f32).sqrt());
    g.mul(acc, inv_sqrt)
}

fn build_atten2_map(
    g: &mut Graph,
    cfg: &RepformersConfig,
    prefix: &str,
    g2: NodeId,
    h2: NodeId,
    nlist_mask: NodeId,
    sw: NodeId,
    nf: usize,
    nloc: usize,
    nnei: usize,
) -> Result<NodeId> {
    let nh = cfg.attn2_nhead;
    let nd = cfg.attn2_hidden;
    let ng2 = cfg.g2_dim;

    // mapqk: [ng2 → 2 * nh * nd] linear, no bias
    let w = g.param(
        format!("{prefix}.mapqk.w"),
        Shape::new(&[ng2, 2 * nh * nd], DType::F32),
    );
    let qk = g.mm(g2, w); // [nf, nloc, nnei, 2*nh*nd]
    // reshape to [nf, nloc, nnei, nd, 2*nh], transpose to [nf, nloc, 2*nh, nnei, nd]
    let qk_reshape_shape = Shape::new(&[nf, nloc, nnei, nd, 2 * nh], DType::F32);
    let qk_r = g.reshape(
        qk,
        vec![nf as i64, nloc as i64, nnei as i64, nd as i64, (2 * nh) as i64],
        qk_reshape_shape,
    );
    let qk_t = g.transpose_(qk_r, vec![0, 1, 4, 2, 3]); // [nf, nloc, 2nh, nnei, nd]
    // split heads
    let q = g.narrow_(qk_t, 2, 0, nh); // [nf, nloc, nh, nnei, nd]
    let k = g.narrow_(qk_t, 2, nh, nh);
    let k_t = g.transpose_(k, vec![0, 1, 2, 4, 3]); // [nf, nloc, nh, nd, nnei]
    let mut attnw = g.mm(q, k_t); // [nf, nloc, nh, nnei, nnei]
    let scale = scalar_const(g, 1.0 / (nd as f32).sqrt());
    attnw = g.mul(attnw, scale);

    if cfg.attn2_has_gate {
        let h2_t = g.transpose_(h2, vec![0, 1, 3, 2]); // [nf, nloc, 3, nnei]
        let gate = g.mm(h2, h2_t); // [nf, nloc, nnei, nnei]
        let gate_5d_shape = Shape::new(&[nf, nloc, 1, nnei, nnei], DType::F32);
        let gate_5d = g.reshape(
            gate,
            vec![nf as i64, nloc as i64, 1, nnei as i64, nnei as i64],
            gate_5d_shape,
        );
        attnw = g.mul(attnw, gate_5d);
    }

    if cfg.smooth {
        // `(attnw + 20) * sw_q * sw_k - 20` rewritten to avoid the
        // rlx-cpu buffer alias on `attnw + 20` (see DPA-1 fix).
        let sw_b = g.reshape(
            sw,
            vec![nf as i64, nloc as i64, 1, 1, nnei as i64],
            Shape::new(&[nf, nloc, 1, 1, nnei], DType::F32),
        );
        let sw_a = g.reshape(
            sw,
            vec![nf as i64, nloc as i64, 1, nnei as i64, 1],
            Shape::new(&[nf, nloc, 1, nnei, 1], DType::F32),
        );
        let sw_prod = g.mul(sw_a, sw_b);
        let aw_scaled = g.mul(attnw, sw_prod);
        let one = scalar_const(g, 1.0);
        let sw_prod_shape = g.shape(sw_prod).clone();
        let sw_minus_one = g.binary(BinaryOp::Sub, sw_prod, one, sw_prod_shape);
        let twenty = scalar_const(g, 20.0);
        let bias_shape = g.shape(sw_minus_one).clone();
        let bias = g.binary(BinaryOp::Mul, twenty, sw_minus_one, bias_shape);
        attnw = g.add(aw_scaled, bias);
    } else {
        // hard mask using -inf where nlist_mask == 0
        let mask_5d = g.reshape(
            nlist_mask,
            vec![nf as i64, nloc as i64, 1, 1, nnei as i64],
            Shape::new(&[nf, nloc, 1, 1, nnei], DType::F32),
        );
        let one = scalar_const(g, 1.0);
        let inv = g.sub(one, mask_5d);
        let neg_inf = scalar_const(g, -1e30);
        let add = g.mul(inv, neg_inf);
        attnw = g.add(attnw, add);
    }
    let aw_shape = g.shape(attnw).clone();
    attnw = g.softmax(attnw, -1, aw_shape);

    // Python zeros both columns (key-masked) and rows (query-masked) post-softmax.
    let mask_k_5d = g.reshape(
        nlist_mask,
        vec![nf as i64, nloc as i64, 1, 1, nnei as i64],
        Shape::new(&[nf, nloc, 1, 1, nnei], DType::F32),
    );
    attnw = g.mul(attnw, mask_k_5d);
    let mask_q_5d = g.reshape(
        nlist_mask,
        vec![nf as i64, nloc as i64, 1, nnei as i64, 1],
        Shape::new(&[nf, nloc, 1, nnei, 1], DType::F32),
    );
    attnw = g.mul(attnw, mask_q_5d);
    if cfg.smooth {
        let sw_b = g.reshape(
            sw,
            vec![nf as i64, nloc as i64, 1, 1, nnei as i64],
            Shape::new(&[nf, nloc, 1, 1, nnei], DType::F32),
        );
        let sw_a = g.reshape(
            sw,
            vec![nf as i64, nloc as i64, 1, nnei as i64, 1],
            Shape::new(&[nf, nloc, 1, nnei, 1], DType::F32),
        );
        attnw = g.mul(attnw, sw_a);
        attnw = g.mul(attnw, sw_b);
    }

    // h2 · h2^T  / √3   → [nf, nloc, nnei, nnei]
    let h2_t = g.transpose_(h2, vec![0, 1, 3, 2]);
    let h2h2 = g.mm(h2, h2_t);
    let inv_sqrt3 = scalar_const(g, 1.0 / 3f32.sqrt());
    let h2h2 = g.mul(h2h2, inv_sqrt3);
    let h2h2_5d_shape = Shape::new(&[nf, nloc, 1, nnei, nnei], DType::F32);
    let h2h2_5d = g.reshape(
        h2h2,
        vec![nf as i64, nloc as i64, 1, nnei as i64, nnei as i64],
        h2h2_5d_shape,
    );
    let aag = g.mul(attnw, h2h2_5d); // [nf, nloc, nh, nnei, nnei]
    // Transpose to [nf, nloc, nnei, nnei, nh]
    Ok(g.transpose_(aag, vec![0, 1, 3, 4, 2]))
}

fn build_atten2_mh_apply(
    g: &mut Graph,
    cfg: &RepformersConfig,
    prefix: &str,
    aag: NodeId, // [nf, nloc, nnei, nnei, nh]
    g2: NodeId,  // [nf, nloc, nnei, ng2]
    nf: usize,
    nloc: usize,
    nnei: usize,
    ng2: usize,
) -> NodeId {
    let nh = cfg.attn2_nhead;
    // Mirror `Atten2MultiHeadApply.call`:
    //   g2v = mapv(g2)                              # [nf, nloc, nnei, ng2*nh]
    //   g2v = reshape(g2v, [nf, nloc, nnei, ng2, nh])
    //   g2v = transpose_01423(g2v)                  # [nf, nloc, nh, nnei, ng2]
    //   AA  = transpose_01423(AA)                   # [nf, nloc, nh, nnei, nnei]
    //   ret = matmul(AA, g2v)                       # [nf, nloc, nh, nnei, ng2]
    //   ret = transpose_01342(ret)                  # [nf, nloc, nnei, ng2, nh]
    //   ret = reshape(ret, [nf, nloc, nnei, ng2*nh])
    //   return head_map(ret)                        # NativeLayer(ng2*nh, ng2)
    let mapv_w = g.param(
        format!("{prefix}.mapv.w"),
        Shape::new(&[ng2, nh * ng2], DType::F32),
    );
    let v = g.mm(g2, mapv_w); // [nf, nloc, nnei, ng2*nh]
    let v_5d_shape = Shape::new(&[nf, nloc, nnei, ng2, nh], DType::F32);
    let v_5d = g.reshape(
        v,
        vec![nf as i64, nloc as i64, nnei as i64, ng2 as i64, nh as i64],
        v_5d_shape,
    );
    // transpose_01423: [0,1,4,2,3] → moves nh from axis 4 to axis 2
    let v_perm = g.transpose_(v_5d, vec![0, 1, 4, 2, 3]); // [nf, nloc, nh, nnei, ng2]
    let aag_perm = g.transpose_(aag, vec![0, 1, 4, 2, 3]); // [nf, nloc, nh, nnei, nnei]
    let out_5d = g.mm(aag_perm, v_perm); // [nf, nloc, nh, nnei, ng2]
    // transpose_01342: [0,1,3,4,2] → moves nh from axis 2 to axis 4
    let out_perm = g.transpose_(out_5d, vec![0, 1, 3, 4, 2]); // [nf, nloc, nnei, ng2, nh]
    let out_flat_shape = Shape::new(&[nf, nloc, nnei, nh * ng2], DType::F32);
    let out_flat = g.reshape(
        out_perm,
        vec![nf as i64, nloc as i64, nnei as i64, (nh * ng2) as i64],
        out_flat_shape,
    );
    // head_map: NativeLayer(ng2 * nh, ng2) — with bias
    let head_w = g.param(
        format!("{prefix}.head_map.w"),
        Shape::new(&[nh * ng2, ng2], DType::F32),
    );
    let head_b = g.param(
        format!("{prefix}.head_map.b"),
        Shape::new(&[ng2], DType::F32),
    );
    let mm = g.mm(out_flat, head_w);
    g.add(mm, head_b)
}

fn build_local_atten(
    g: &mut Graph,
    cfg: &RepformersConfig,
    prefix: &str,
    g1: NodeId,    // [nf, nloc, ng1]
    gg1: NodeId,   // [nf, nloc, nnei, ng1]
    nlist_mask: NodeId,
    sw: NodeId,
    nf: usize,
    nloc: usize,
    nnei: usize,
    ng1: usize,
) -> Result<NodeId> {
    // Mirror `LocalAtten.call` exactly:
    //   g1q = mapq(g1)              [nf, nloc, nd*nh]  → reshape [nf, nloc, nd, nh] → transpose → [nf, nloc, nh, nd]
    //   gg1kv = mapkv(gg1)          [nf, nloc, nnei, (nd+ng1)*nh] → reshape [nf, nloc, nnei, nd+ng1, nh] → transpose_01423 → [nf, nloc, nh, nnei, nd+ng1]
    //   gg1k = gg1kv[..., :nd]      [nf, nloc, nh, nnei, nd]
    //   gg1v = gg1kv[..., nd:]      [nf, nloc, nh, nnei, ng1]
    //   attnw = (g1q[None] @ gg1k^T) / √nd
    //   attnw = squeeze(axis=-2)    [nf, nloc, nh, nnei]
    //   smooth: attnw = (attnw + shift) * sw - shift ; then ×sw post-softmax
    //   ret = (attnw[None] @ gg1v) → [nf, nloc, nh, ng1] → reshape [nf, nloc, nh*ng1] → head_map → [nf, nloc, ng1]
    let nh = cfg.attn1_nhead;
    let nd = cfg.attn1_hidden;

    // mapq: NativeLayer(ng1, nd*nh, bias=False)
    let mapq_w = g.param(
        format!("{prefix}.mapq.w"),
        Shape::new(&[ng1, nd * nh], DType::F32),
    );
    let g1q = g.mm(g1, mapq_w); // [nf, nloc, nd*nh]
    // reshape to [nf, nloc, nd, nh] then matrix_transpose (swap last two) → [nf, nloc, nh, nd]
    let g1q_4d = g.reshape(
        g1q,
        vec![nf as i64, nloc as i64, nd as i64, nh as i64],
        Shape::new(&[nf, nloc, nd, nh], DType::F32),
    );
    let g1q_perm = g.transpose_(g1q_4d, vec![0, 1, 3, 2]); // [nf, nloc, nh, nd]

    // mapkv: NativeLayer(ng1, (nd+ng1)*nh, bias=False)
    let mapkv_w = g.param(
        format!("{prefix}.mapkv.w"),
        Shape::new(&[ng1, (nd + ng1) * nh], DType::F32),
    );
    let gg1kv = g.mm(gg1, mapkv_w); // [nf, nloc, nnei, (nd+ng1)*nh]
    let gg1kv_5d = g.reshape(
        gg1kv,
        vec![nf as i64, nloc as i64, nnei as i64, (nd + ng1) as i64, nh as i64],
        Shape::new(&[nf, nloc, nnei, nd + ng1, nh], DType::F32),
    );
    // transpose_01423: [0,1,4,2,3] → [nf, nloc, nh, nnei, nd+ng1]
    let gg1kv_perm = g.transpose_(gg1kv_5d, vec![0, 1, 4, 2, 3]);
    let gg1k = g.narrow_(gg1kv_perm, 4, 0, nd);   // [nf, nloc, nh, nnei, nd]
    let gg1v = g.narrow_(gg1kv_perm, 4, nd, ng1); // [nf, nloc, nh, nnei, ng1]

    // attnw = (g1q[..., None, :] @ gg1k^T) / √nd → [nf, nloc, nh, 1, nnei] → squeeze
    let g1q_5d = g.reshape(
        g1q_perm,
        vec![nf as i64, nloc as i64, nh as i64, 1, nd as i64],
        Shape::new(&[nf, nloc, nh, 1, nd], DType::F32),
    );
    let gg1k_t = g.transpose_(gg1k, vec![0, 1, 2, 4, 3]); // [nf, nloc, nh, nd, nnei]
    let mut attnw = g.mm(g1q_5d, gg1k_t); // [nf, nloc, nh, 1, nnei]
    let scale = scalar_const(g, 1.0 / (nd as f32).sqrt());
    attnw = g.mul(attnw, scale);
    // squeeze axis=-2 → [nf, nloc, nh, nnei]
    let attnw_4d_shape = Shape::new(&[nf, nloc, nh, nnei], DType::F32);
    let mut attnw = g.reshape(
        attnw,
        vec![nf as i64, nloc as i64, nh as i64, nnei as i64],
        attnw_4d_shape,
    );
    if cfg.smooth {
        // `(attnw + 20) * sw - 20` rewritten as `attnw * sw + 20 * (sw - 1)`
        let sw_4d = g.reshape(
            sw,
            vec![nf as i64, nloc as i64, 1, nnei as i64],
            Shape::new(&[nf, nloc, 1, nnei], DType::F32),
        );
        let aw_scaled = g.mul(attnw, sw_4d);
        let one = scalar_const(g, 1.0);
        let sw_shape = g.shape(sw_4d).clone();
        let sw_minus_one = g.binary(BinaryOp::Sub, sw_4d, one, sw_shape);
        let twenty = scalar_const(g, 20.0);
        let bias_shape = g.shape(sw_minus_one).clone();
        let bias = g.binary(BinaryOp::Mul, twenty, sw_minus_one, bias_shape);
        attnw = g.add(aw_scaled, bias);
    } else {
        let mask_4d = g.reshape(
            nlist_mask,
            vec![nf as i64, nloc as i64, 1, nnei as i64],
            Shape::new(&[nf, nloc, 1, nnei], DType::F32),
        );
        let one = scalar_const(g, 1.0);
        let inv = g.sub(one, mask_4d);
        let neg_inf = scalar_const(g, -1e30);
        let add = g.mul(inv, neg_inf);
        attnw = g.add(attnw, add);
    }
    let aw_shape = g.shape(attnw).clone();
    attnw = g.softmax(attnw, -1, aw_shape);
    // Zero out masked queries
    let mask_4d = g.reshape(
        nlist_mask,
        vec![nf as i64, nloc as i64, 1, nnei as i64],
        Shape::new(&[nf, nloc, 1, nnei], DType::F32),
    );
    attnw = g.mul(attnw, mask_4d);
    if cfg.smooth {
        let sw_4d = g.reshape(
            sw,
            vec![nf as i64, nloc as i64, 1, nnei as i64],
            Shape::new(&[nf, nloc, 1, nnei], DType::F32),
        );
        attnw = g.mul(attnw, sw_4d);
    }

    // ret = attnw[..., None, :] @ gg1v → [nf, nloc, nh, 1, ng1]
    let attnw_5d = g.reshape(
        attnw,
        vec![nf as i64, nloc as i64, nh as i64, 1, nnei as i64],
        Shape::new(&[nf, nloc, nh, 1, nnei], DType::F32),
    );
    let ret = g.mm(attnw_5d, gg1v); // [nf, nloc, nh, 1, ng1]
    // squeeze axis=-2 → [nf, nloc, nh, ng1], reshape to [nf, nloc, nh*ng1]
    let ret_flat_shape = Shape::new(&[nf, nloc, nh * ng1], DType::F32);
    let ret_flat = g.reshape(
        ret,
        vec![nf as i64, nloc as i64, (nh * ng1) as i64],
        ret_flat_shape,
    );
    // head_map: NativeLayer(nh*ng1, ng1) — with bias
    let head_w = g.param(
        format!("{prefix}.head_map.w"),
        Shape::new(&[nh * ng1, ng1], DType::F32),
    );
    let head_b = g.param(
        format!("{prefix}.head_map.b"),
        Shape::new(&[ng1], DType::F32),
    );
    let mm = g.mm(ret_flat, head_w);
    Ok(g.add(mm, head_b))
}

// Keep activation reference live across optional code paths.
#[allow(dead_code)]
fn _keep_act_path(_a: Activation, _b: BinaryOp) {}

#[cfg(test)]
mod tests {
    use super::*;

    #[test]
    fn repformers_builds_default() {
        let cfg = RepformersConfig {
            rcut: 6.0,
            rcut_smth: 0.5,
            sel: 16,
            ntypes: 2,
            nlayers: 1,
            g1_dim: 16,
            g2_dim: 8,
            axis_neuron: 2,
            update_g1_has_conv: true,
            update_g1_has_drrd: true,
            update_g1_has_grrg: true,
            update_g1_has_attn: true,
            update_g2_has_g1g1: true,
            update_g2_has_attn: true,
            update_h2: false,
            attn1_hidden: 8,
            attn1_nhead: 2,
            attn2_hidden: 4,
            attn2_nhead: 2,
            attn2_has_gate: false,
            activation_function: "tanh".into(),
            update_style: "res_avg".into(),
            smooth: true,
            use_sqrt_nnei: true,
            g1_out_conv: true,
            g1_out_mlp: true,
            ln_eps: 1e-5,
            epsilon: 1e-4,
        };
        let mut g = Graph::new("repformers");
        let nf = 1;
        let nloc = 2;
        let nall = nloc;
        let nnei = cfg.sel;
        let g1_ext = g.input("g1_ext", Shape::new(&[nf, nall, cfg.g1_dim], DType::F32));
        let env_mat = g.input("em", Shape::new(&[nf, nloc, nnei, 4], DType::F32));
        let nlist = g.input("nl", Shape::new(&[nf, nloc, nnei], DType::I32));
        let nlist_mask = g.input("nm", Shape::new(&[nf, nloc, nnei], DType::F32));
        let sw = g.input("sw", Shape::new(&[nf, nloc, nnei], DType::F32));
        let atype = g.input("at", Shape::new(&[nf, nloc], DType::I32));
        let out = build_repformers(
            &mut g,
            &cfg,
            RepformersInputs {
                g1_ext,
                env_mat_raw: env_mat,
                nlist,
                nlist_mask,
                sw,
                atype_loc: atype,
                mapping: None,
            },
            nf,
            nloc,
            nall,
        )
        .expect("build");
        let s = g.shape(out.g1);
        assert_eq!(s.dim(2), rlx_ir::Dim::Static(cfg.g1_dim));
    }
}