diff --git a/src/Auto3D/ASE/thermo.py b/src/Auto3D/ASE/thermo.py index 6b1a432..f8dff9a 100644 --- a/src/Auto3D/ASE/thermo.py +++ b/src/Auto3D/ASE/thermo.py @@ -1516,10 +1516,14 @@ def _load_hessian_model(model_name: str, device) -> ModelAdapter: # compile_model=False: torch.compile guards on dtype, and nothing in # this autograd-Hessian path benefits from it anyway. adapter = create_model(model_name, device, compile_model=False, use_cache=False) - # In place, and on the adapter's own module, exactly as before: the - # adapter is what gets returned now, but the fp64 tensor it will feed the - # model is the same one, so no reported frequency moves. - adapter.model.double() + # In place, through the contract rather than past it. This was + # `adapter.model.double()`, which reached the module only + # BaseModelAdapter happens to store -- so an otherwise conforming + # structural adapter raised AttributeError here, and mypy's report of it + # was one of the errors `|| true` discarded. The operation underneath is + # unchanged (see BaseModelAdapter.to_double), so no reported frequency + # moves. + adapter.to_double() return adapter # AIMNET or any aimnet registry alias: ModelFactory resolves the "AIMNET" # legacy alias to the registry default internally (see diff --git a/src/Auto3D/models/adapter.py b/src/Auto3D/models/adapter.py index 48a3dc6..7e0f25e 100644 --- a/src/Auto3D/models/adapter.py +++ b/src/Auto3D/models/adapter.py @@ -264,6 +264,20 @@ def analytic_hessian( """ return None + def to_double(self) -> None: + """Promote the wrapped module's weights to float64, in place. + + ``self.model.double()`` rather than ``self.double()``, although this + class is itself an ``nn.Module`` and the two coincide for every adapter + that registers no other child. They are not guaranteed to: ``Module.double`` + recurses into every registered submodule, so a subclass that holds a second + one would have it upcast too. This is the exact operation + ``Auto3D.ASE.thermo`` performed before the call moved onto the contract, + and keeping it exact is what makes "no reported frequency moves" a claim + rather than a hope. + """ + self.model.double() + @abstractmethod def forward( self, @@ -373,6 +387,32 @@ def analytic_hessian( ) return result["hessian"] + def to_double(self) -> None: + """Refuse: AIMNet2 has no meaningful whole-graph fp64 form. + + Two independent reasons, either sufficient. Whole-graph fp64 through + AIMNet2 is false precision -- the network is trained and evaluated in + fp32 -- and this adapter never needs the upcast anyway, because + :meth:`analytic_hessian` above means it is never differentiated by + autograd. ``Auto3D.ASE.thermo._load_hessian_model`` accordingly upcasts + the ANI and custom-model branches and not this one. + + The mechanism would also be wrong, not merely unnecessary: ``self.model`` + is ``self._calc.model``, the same object, so upcasting it here mutates the + module underneath an ``AIMNet2Calculator`` that prepares its own fp32 + input tensors. + + Inheriting :meth:`BaseModelAdapter.to_double` would make all of that a + silent, working-looking call. This raise is the same judgment + :meth:`analytic_hessian` documents for ``None``: a member that cannot + honestly do what it says must say so rather than appear to comply. + """ + raise NotImplementedError( + "AIMNet2 has no fp64 form: it is trained and evaluated in float32, " + "and its Hessian is analytic, so it is never differentiated by " + "autograd and never needs the upcast. Use analytic_hessian instead." + ) + def energy( self, coords: torch.Tensor, diff --git a/src/Auto3D/models/contract.py b/src/Auto3D/models/contract.py index 0b53e95..890df5f 100644 --- a/src/Auto3D/models/contract.py +++ b/src/Auto3D/models/contract.py @@ -337,6 +337,37 @@ def analytic_hessian( """ ... + def to_double(self) -> None: + """Promote this model's weights to float64, in place. + + The counterpart to :meth:`analytic_hessian` returning ``None``. An + adapter with no native second derivative is differentiated by + ``torch.autograd.functional.hessian`` on the fp64 geometry + ``Auto3D.ASE.thermo.vib_hessian`` builds, and fp32 weights make that fp64 + request meaningless -- :meth:`energy` is dtype-*preserving* precisely so + the caller can choose, but choosing fp64 is only real once the weights + follow. This member is how the caller says so. + + It exists because the caller previously said it as + ``adapter.model.double()``, reaching past this contract to an attribute + only :class:`~Auto3D.models.adapter.BaseModelAdapter` happens to define. + Every structural adapter -- test doubles, and anything a downstream user + writes -- has no ``.model``, so the one path that upcasts raised + ``AttributeError`` on an otherwise conforming object. + + In place, and returning nothing: the sole caller discards the result, and + ``nn.Module.double()``'s return-self convention would imply an + out-of-place option that does not exist. + + Raises: + NotImplementedError: The model has no meaningful fp64 form. Raising + is the honest answer for a backend whose energy pipeline is fp32 + by design -- see :meth:`Auto3D.models.adapter.AIMNet2Adapter.to_double`. + An implementation must not make this a silent no-op: the caller + would then differentiate fp32 weights while believing otherwise. + """ + ... + def missing_adapter_members(obj: Any) -> list[str]: """Members :class:`ModelAdapter` requires that ``obj`` does not provide. diff --git a/tests/helpers_adapter.py b/tests/helpers_adapter.py index 37981f2..4407bb0 100644 --- a/tests/helpers_adapter.py +++ b/tests/helpers_adapter.py @@ -84,6 +84,16 @@ def analytic_hessian(self, coords, species, charges): self.calls.append({"dtype": coords.dtype, "kind": "analytic_hessian"}) return self._hessian + def to_double(self) -> None: + """Recorded, not performed: there are no weights here to upcast. + + This double computes its energies from ``coords`` alone, so it is already + dtype-preserving and an upcast has nothing to act on. Recording the call + still matters -- it is how a test checks that the fp64 request reached the + adapter at all. + """ + self.calls.append({"kind": "to_double"}) + class AdapterModuleMixin: """Makes an ``nn.Module`` test double satisfy ``ModelAdapter``. @@ -114,6 +124,16 @@ def analytic_hessian(self, coords, species, charges): """No native second derivative -- ``BaseModelAdapter``'s own default.""" return None + def to_double(self) -> None: + """Upcast whatever this double registered, matching the real adapters. + + These doubles ARE ``nn.Module``s, so unlike :class:`FakeAdapter` there may + genuinely be parameters to promote. ``self.double()`` rather than + ``self.model.double()`` because a mixed-in double need not wrap a module + at all -- several declare their forward inline and hold no ``.model``. + """ + self.double() + def padded_batch(n_mols: int = 2, n_atoms: int = 3): """Tensors shaped like :func:`Auto3D.batch_opt.padding.pad_from_mols`.""" diff --git a/tests/test_adapter_fp64_contract.py b/tests/test_adapter_fp64_contract.py new file mode 100644 index 0000000..9a21a7f --- /dev/null +++ b/tests/test_adapter_fp64_contract.py @@ -0,0 +1,118 @@ +# tests/test_adapter_fp64_contract.py +"""``to_double``: the fp64 upcast expressed on the contract, not through it. + +``Auto3D.ASE.thermo._load_hessian_model`` used to write ``adapter.model.double()`` +-- reaching past :class:`~Auto3D.models.contract.ModelAdapter` to the module the +adapter happens to wrap. It worked only because every in-tree adapter derives from +:class:`~Auto3D.models.adapter.BaseModelAdapter`, which stores one; a conforming +structural adapter (which production has always accepted -- see +``tests/helpers_adapter``) has no ``.model`` at all and died with ``AttributeError`` +inside the thermochemistry path. mypy said so on every run and ``|| true`` +discarded it. +""" + +from __future__ import annotations + +import torch +from torch import nn + +from Auto3D.models.adapter import BaseModelAdapter +from Auto3D.models.contract import missing_adapter_members +from tests.helpers_adapter import FakeAdapter + + +def test_to_double_is_part_of_the_adapter_contract(): + """The upcast is a member, so the gate rejects an adapter that lacks it. + + ``missing_adapter_members`` derives from the Protocol, so this asserts the + member was added to :class:`ModelAdapter` itself rather than only to the + implementation base class -- which is the whole distinction the contract + module exists to keep. + """ + + class _WithoutUpcast: + coord_pad = 0.0 + species_pad = -1 + + def to_species(self, numbers): + return list(numbers) + + def forward(self, coords, species, charges, atom_mask=None): + raise AssertionError("not called") + + def energy(self, coords, species, charges, atom_mask=None): + raise AssertionError("not called") + + def analytic_hessian(self, coords, species, charges): + return None + + assert missing_adapter_members(_WithoutUpcast()) == ["to_double"] + + +def test_base_adapter_to_double_upcasts_the_wrapped_module(): + """The default implementation is exactly the operation it replaced. + + ``self.model.double()`` and not ``self.double()``: the adapter is itself an + ``nn.Module``, so the second would additionally recurse into anything else a + subclass registered as a child -- for ``AIMNet2Adapter`` that reaches modules + the old call never touched. Byte-for-byte the previous operation is the point. + """ + + class _Concrete(BaseModelAdapter): + def forward(self, coords, species, charges, atom_mask=None): + raise AssertionError("not called") + + adapter = _Concrete(nn.Linear(2, 2), torch.device("cpu")) + assert next(adapter.model.parameters()).dtype == torch.float32 + + adapter.to_double() + + assert next(adapter.model.parameters()).dtype == torch.float64 + + +class TestLoadHessianModelUpcastsThroughTheContract: + """``_load_hessian_model`` must not reach past the adapter it was handed. + + ``create_model`` is monkeypatched in both tests, so no NNP is loaded and + torchani need not be installed; only which contract member the branch calls + is under test. + """ + + def _install(self, monkeypatch): + from Auto3D.ASE import thermo as thermo_mod + + calls: list[str] = [] + + class _NoDotModel(FakeAdapter): + """Conforming, and deliberately without a ``.model`` attribute. + + That absence IS the assertion: the old ``adapter.model.double()`` + raises ``AttributeError`` here rather than quietly passing. + """ + + def to_double(self): + calls.append("to_double") + + monkeypatch.setattr(thermo_mod, "create_model", lambda *a, **k: _NoDotModel()) + return calls, thermo_mod + + def test_ani_branch_upcasts_through_to_double(self, monkeypatch): + calls, thermo_mod = self._install(monkeypatch) + + thermo_mod._load_hessian_model("ANI2xt", torch.device("cpu")) + + assert calls == ["to_double"] + + def test_aimnet_branch_does_not_upcast_at_all(self, monkeypatch): + """Whole-graph fp64 through AIMNet2 is false precision, and its Hessian + is analytic anyway -- so this branch must leave the model fp32. + + Guards the direction the fast tests could not reach before: with + ``create_model`` patched, nothing is loaded, so the branch that must + *not* upcast is finally checkable without a real NNP. + """ + calls, thermo_mod = self._install(monkeypatch) + + thermo_mod._load_hessian_model("AIMNET", torch.device("cpu")) + + assert calls == [] diff --git a/tests/test_batchopt.py b/tests/test_batchopt.py index b0b9060..2f8af2b 100644 --- a/tests/test_batchopt.py +++ b/tests/test_batchopt.py @@ -357,23 +357,18 @@ def test_charges_reach_the_model_as_float32(): seen_dtypes = [] - class _RecordingAdapter: - coord_pad = 0.0 - species_pad = -1 + class _RecordingAdapter(FakeAdapter): + """Only the dtype recording is this test's own. + + Was a hand-rolled duplicate of every ``FakeAdapter`` member; that is the + shape ``tests/helpers_adapter`` exists to absorb, and it went red the + first time the adapter contract gained a member. + """ def forward(self, coords, species, charges, atom_mask=None): seen_dtypes.append(charges.dtype) return torch.zeros(coords.shape[0]), torch.zeros_like(coords) - def to_species(self, numbers): - return numbers - - def energy(self, coords, species, charges, atom_mask=None): - return torch.zeros(coords.shape[0]) - - def analytic_hessian(self, coords, species, charges): - return None - model = EnForce_ANI(_RecordingAdapter(), batchsize_atoms=1024) coord = torch.randn(2, 3, 3) numbers = torch.tensor([[6, 1, 1], [6, 1, 1]], dtype=torch.long)