[Code scan] Atomic virial outputs squeeze the wrong dimension

[njzjz]

This issue was found by a Codex global scan of the repository at commit 19f9265.

energy_derv_c is produced with shape (nf, nall, 1, 9), but the public atom_virial / extended_virial paths call .squeeze(-3). For nloc > 1, -3 is the atom axis, so the singleton derivative axis remains and the output stays 4D.

deepmd-gnn/deepmd_gnn/mace.py

Lines 833 to 836 in 19f9265

	model_predict["atom_virial"] = model_ret_lower["energy_derv_c"][
	:,
	:nloc,
	].squeeze(-3)

deepmd-gnn/deepmd_gnn/mace.py

Lines 931 to 934 in 19f9265

	model_predict["extended_force"] = model_ret["energy_derv_r"].squeeze(-2)
	model_predict["virial"] = model_ret["energy_derv_c_redu"].squeeze(-2)
	if do_atomic_virial:
	model_predict["extended_virial"] = model_ret["energy_derv_c"].squeeze(-3)

deepmd-gnn/deepmd_gnn/nequip.py

Lines 559 to 562 in 19f9265

	model_predict["atom_virial"] = model_ret_lower["energy_derv_c"][
	:,
	:nloc,
	].squeeze(-3)

deepmd-gnn/deepmd_gnn/nequip.py

Lines 628 to 631 in 19f9265

	model_predict["extended_force"] = model_ret["energy_derv_r"].squeeze(-2)
	model_predict["virial"] = model_ret["energy_derv_c_redu"].squeeze(-2)
	if do_atomic_virial:
	model_predict["extended_virial"] = model_ret["energy_derv_c"].squeeze(-3)

The translated output metadata has the same squeeze target:

deepmd-gnn/deepmd_gnn/mace.py

Lines 1791 to 1796 in 19f9265

	output_def["force"] = deepcopy(out_def_data["energy_derv_r"])
	output_def["force"].squeeze(-2)
	output_def["virial"] = deepcopy(out_def_data["energy_derv_c_redu"])
	output_def["virial"].squeeze(-2)
	output_def["atom_virial"] = deepcopy(out_def_data["energy_derv_c"])
	output_def["atom_virial"].squeeze(-3)

deepmd-gnn/deepmd_gnn/nequip.py

Lines 1102 to 1107 in 19f9265

	output_def["force"] = deepcopy(out_def_data["energy_derv_r"])
	output_def["force"].squeeze(-2)
	output_def["virial"] = deepcopy(out_def_data["energy_derv_c_redu"])
	output_def["virial"].squeeze(-2)
	output_def["atom_virial"] = deepcopy(out_def_data["energy_derv_c"])
	output_def["atom_virial"].squeeze(-3)

I verified the metadata path currently reports atom_virial as [1, 9], and the runtime tensor path can return (nf, nloc, 1, 9) instead of (nf, nloc, 9).

Suggested fix: change these squeezes to .squeeze(-2) and add shape assertions for MACE and NequIP when do_atomic_virial=True.