[BUG] DPA4C compression crashes for trainable=False descriptors

[njzjz-bot]

Summary

DescrptDPA4C.enable_compression() crashes when the descriptor was constructed with the supported trainable=False option.

This reproduces on origin/master at 8cfd46e37448.

Reproduction

from deepmd.pt_expt.descriptor.dpa4c import DescrptDPA4C

d = DescrptDPA4C(
    rcut=3.0,
    ntypes=2,
    channels=8,
    lmax=2,
    n_radial=4,
    precision="float32",
    trainable=False,
).eval()
d.enable_compression(0.5, table_stride_1=0.2)

Actual result:

deepmd/kernels/cuda/dpa4c/graph_compress.py:525
sample_parameter = next(descriptor.radial_embedding.parameters())
StopIteration

Cause

build_radial_table() obtains the device and dtype from the first radial-embedding parameter. With trainable=False, the frozen radial MLP weights are registered as buffers, so radial_embedding.parameters() is empty.

Expected behavior

Compression should work for a frozen descriptor. trainable changes parameter registration, not the learned function or compression eligibility.

The device/dtype should be obtained from a tensor that exists in both trainable and frozen forms, and a regression test should enable compression on a trainable=False DPA4C descriptor.

---

Coding agent: Codex
Codex version: codex-cli 0.149.0
Model: gpt-5.6-sol
Reasoning effort: xhigh

[dosubot]

Your diagnosis is correct. Looking at the source at commit 8cfd46e, there are actually four call sites in graph_compress.py that use next(....parameters()) and will crash with StopIteration when trainable=False turns weights into buffers:

1. build_radial_table — next(descriptor.radial_embedding.parameters()) to infer device [1]
2. build_compression_artifacts — next(descriptor.parameters()) to check dtype and get device [2]
3. charge_state_artifacts — next(descriptor.parameters()).device [3]
4. build_charge_state_artifacts — next(descriptor.parameters()).device when building the charge-state tensor [4]

All four need the same fix: fall back to buffers when the parameter iterator is empty. A clean pattern is:

import itertools

# In build_radial_table:
sample_tensor = next(itertools.chain(
    descriptor.radial_embedding.parameters(),
    descriptor.radial_embedding.buffers(),
))

# In the other three sites:
sample_tensor = next(itertools.chain(
    descriptor.parameters(),
    descriptor.buffers(),
))

The trainable attribute already has an inconsistent-default issue tracked in the capability-probe cleanup, and a related buffer-vs-parameter question was discussed for Paddle DPA3 (though that one turned out to be a non-issue because Paddle aliases requires_grad to stop_gradient). This DPA4C case is a genuine crash, since PyTorch's register_buffer removes the tensor from parameters() entirely.

A regression test calling enable_compression on a trainable=False descriptor and asserting no StopIteration would prevent this from recurring.

To reply, just mention @dosu.

---

Share context across your team and agents. Try Dosu.