[daylight-00]

Summary

• ESMFold2InputBuilder.fold() accepts precomputed lm_hidden_states (e.g. to reuse one ESMC pass across seeds or sampler settings).
• Supplied states are checked against the input and the live LM shim: batch/tokens, layer count, width, device, floating dtype. That they were computed from this input is the caller's responsibility (stated in the docstring). Passing lm_mask_pct alongside them raises (the mask only acts inside the ESMC pass).
• New public compute_lm_hidden_states() on EsmFold2Model and EsmFold2ExperimentalModel: wraps the existing private helper and resolves lm_mask_pct=None from config. On EsmFold2Model, a standalone call returns the ESMC backbone to CPU when _offload_esmc=True; the experimental model does not implement ESMC offloading.
• fold() without ESMC and without states behaves as before.

Tests

• New CPU tests (21): supplied states reaching the shim, the dtype cast, each validation, the public method matching forward()'s own LM pass (release and experimental), the backbone requirement, and offload / lm_mask_pct defaults through a mock seam.
• tests/models on CPU against main (90d6497): 308 passed / 7 xfailed on main, 329 passed / 7 xfailed here; the 21 new tests pass and no existing test changes status.
• gpu smoke test (tiny model, CUDA): a standalone call leaves ESMC on CPU and its states fold. Passed on an RTX 6000 Ada.
• ruff 0.15.16 passes; ty 0.0.49 reports the same 82 diagnostics on main and on this branch, none new.
• Not run: an fp8 ESMC (the GPU smoke uses a tiny ESMC).

Notes

• Updated onto main at 90d6497 by a merge commit; the tests use its fixtures (ccd_subset, tiny_esmfold2_experimental).
• Earlier revisions also made fold() refuse a model without ESMC. That is left out, to keep this PR to the hidden-states API.

[Copilot]

Copilot review overview

🟡 Changes recommended

Three unresolved moderate findings remain in the reviewed implementation.

Review effort: Lite
Findings: 1

Open (1)

• Floating-point validation permits incompatible state dtypes · New

What changed in this PR

Adds reusable precomputed ESMC hidden-state support to ESMFold2 folding, public standalone LM-state computation, and expanded API tests.

Changes:

• Validates and forwards lm_hidden_states through fold().
• Adds compute_lm_hidden_states() to both model variants.
• Extends tests for validation, offloading, masking, and reuse.

File	Summary	Review findings
tests/​models/​esmfold2_api_test.py	Adds coverage for new LM-state workflows.	None
esm/​models/​esmfold2/​processor.py	Validates and forwards supplied LM states.	Moderate findings: adapter-aware shim lookup (1 vote); floating-dtype contract validation (3 votes).
esm/​models/​esmfold2/​model.py	Adds standalone LM-state computation.	None
esm/​models/​esmfold2/​experimental.py	Adds the experimental model helper.	Moderate autocast precision discrepancy (1 vote).

---

💡 Add a code-review agent skill or configure MCP servers for context-aware, tailored reviews. Learn more in the docs.

[Copilot]

Copilot review overview

🟡 Changes recommended

Moderate issues remain in inference-mode handling, experimental offloading, and adapter-backed shim validation.

Review effort: Lite
Findings: 1

Open (1)

• Restore ESMC to CPU after standalone pass · New

Resolved since last review (1)

• Floating-point validation permits incompatible state dtypes