[Code scan] Respect model precision without leaking torch default dtype

[njzjz]

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

MACE accepts precision through config, but direct construction drops it through **kwargs and builds the MACE network from the ambient torch.get_default_dtype():

deepmd-gnn/deepmd_gnn/mace.py

Lines 270 to 291 in 19f9265

	def __init__(
	self,
	type_map: list[str],
	sel: int,
	r_max: float = 5.0,
	num_radial_basis: int = 8,
	num_cutoff_basis: int = 5,
	max_ell: int = 3,
	interaction: str = "RealAgnosticResidualInteractionBlock",
	num_interactions: int = 2,
	hidden_irreps: str = "128x0e + 128x1o",
	pair_repulsion: bool = False,
	distance_transform: str = "None",
	correlation: int = 3,
	gate: str = "silu",
	MLP_irreps: str = "16x0e",
	radial_type: str = "bessel",
	radial_MLP: list[int] = [64, 64, 64], # noqa: B006
	std: float = 1,
	avg_num_neighbors: float | None = None,
	enable_cueq: bool = False,
	**kwargs: Any, # noqa: ANN401

deepmd-gnn/deepmd_gnn/mace.py

Lines 296 to 315 in 19f9265

	self.params: dict[str, Any] = {
	"type_map": type_map,
	"sel": sel,
	"r_max": r_max,
	"num_radial_basis": num_radial_basis,
	"num_cutoff_basis": num_cutoff_basis,
	"max_ell": max_ell,
	"interaction": interaction,
	"num_interactions": num_interactions,
	"hidden_irreps": hidden_irreps,
	"pair_repulsion": pair_repulsion,
	"distance_transform": distance_transform,
	"correlation": correlation,
	"gate": gate,
	"MLP_irreps": MLP_irreps,
	"radial_type": radial_type,
	"radial_MLP": radial_MLP,
	"std": std,
	"avg_num_neighbors": avg_num_neighbors,
	"enable_cueq": enable_cueq,

deepmd-gnn/deepmd_gnn/mace.py

Lines 338 to 359 in 19f9265

	self.model = _make_mace_network(
	r_max=r_max,
	num_radial_basis=num_radial_basis,
	num_cutoff_basis=num_cutoff_basis,
	max_ell=max_ell,
	interaction=interaction,
	num_interactions=num_interactions,
	num_elements=self.ntypes,
	hidden_irreps=hidden_irreps,
	atomic_numbers=atomic_numbers,
	avg_num_neighbors=self.avg_num_neighbors,
	pair_repulsion=pair_repulsion,
	distance_transform=distance_transform,
	correlation=correlation,
	gate=gate,
	MLP_irreps=MLP_irreps,
	std=std,
	radial_MLP=radial_MLP,
	radial_type=radial_type,
	enable_cueq=enable_cueq and not self._disable_cueq_for_freeze,
	script_model=not self._disable_cueq_for_freeze,
	)

The get_model() config path sets the process-global default dtype and never restores it:

deepmd-gnn/deepmd_gnn/mace.py

Lines 1823 to 1831 in 19f9265

	precision = model_params.pop("precision", "float32")
	if precision == "float32":
	torch.set_default_dtype(torch.float32)
	elif precision == "float64":
	torch.set_default_dtype(torch.float64)
	else:
	msg = f"precision {precision} not supported"
	raise ValueError(msg)
	model = cls(**model_params)

I verified both effects locally:

MaceModel(..., precision="float64") with global float32 builds atomic_energies as torch.float32
MaceModel.get_model(..., precision="float64") leaves torch.get_default_dtype() == torch.float64

NequIP has the complementary failure mode: it passes model_dtype, but does not temporarily raise PyTorch's default dtype before model_from_config(), so precision="float64" can fail when the caller's default dtype is still float32.

deepmd-gnn/deepmd_gnn/nequip.py

Lines 258 to 299 in 19f9265

	precision: str = "float32",
	**kwargs: Any, # noqa: ANN401
	) -> None:
	super().__init__(**kwargs)
	self.params = {
	"type_map": type_map,
	"sel": sel,
	"r_max": r_max,
	"num_layers": num_layers,
	"l_max": l_max,
	"num_features": num_features,
	"nonlinearity_type": nonlinearity_type,
	"parity": parity,
	"num_basis": num_basis,
	"BesselBasis_trainable": BesselBasis_trainable,
	"PolynomialCutoff_p": PolynomialCutoff_p,
	"invariant_layers": invariant_layers,
	"invariant_neurons": invariant_neurons,
	"use_sc": use_sc,
	"irreps_edge_sh": irreps_edge_sh,
	"feature_irreps_hidden": feature_irreps_hidden,
	"chemical_embedding_irreps_out": chemical_embedding_irreps_out,
	"conv_to_output_hidden_irreps_out": conv_to_output_hidden_irreps_out,
	"precision": precision,
	}
	self.type_map = type_map
	self.ntypes = len(type_map)
	self.preset_out_bias: dict[str, list] = {"energy": []}
	self._observed_type = None
	self.mm_types = []
	self.sel = sel
	self.num_layers = num_layers
	for ii, tt in enumerate(type_map):
	if not tt.startswith("m") and tt not in {"HW", "OW"}:
	self.preset_out_bias["energy"].append(None)
	else:
	self.preset_out_bias["energy"].append([0])
	self.mm_types.append(ii)
	self.rcut = r_max
	nequip_model = _make_nequip_network(self.params, self.ntypes)
	self.model = script(nequip_model.to(env.DEVICE))

Local reproduction:

ValueError: Overall default_dtype=float32, but model_dtype=float64 is a higher precision- change default_dtype to float64

Suggested fix: validate precision explicitly in both constructors, build the underlying MACE/NequIP networks under a temporary dtype context, store the effective precision in serialized params, and always restore the previous global default dtype.