Add ESMFold2 tests + sub-config model_types
ESMFold2 is an all-atom structure predictor whose forward takes ~18 structural
feature tensors and returns a plain dict (not a ModelOutput), so it doesn't fit
ModelTesterMixin. Following the sanctioned pattern for such models, register the
test file in check_repo's TEST_FILES_WITH_NO_COMMON_TESTS and provide focused
coverage:
- ESMFold2ConfigTest: ConfigTester common tests (incl. composite sub-config
save/load), `type` validation, nested round-trip, attn-impl propagation.
- ESMFold2ModelTest (CPU): full pure-PyTorch forward via infer_protein with no
ESMC backbone (LM conditioning skipped) under both sdpa and eager; SWA config
dispatch; weight-level save/load fidelity + a usable reloaded model.
- ESMFold2IntegrationTest: @slow real-weight fold on biohub/ESMFold2 (GPU-gated).
To make ConfigTester's composite test pass, each sub-config now declares a
unique `model_type` (e.g. "esmfold2_inputs_embedder") — the CLIP pattern — so
that `SubConfig.from_pretrained(<composite dir>)` extracts the matching nested
dict (configuration_utils keys this off model_type). ESM dodges this test
because its sub-config is None by default; ESMFold2's are always present.
check_repo: `modeling_esmfold2_common` (shared building blocks, no public model)
is added to get_model_modules' _ignore_modules. The remaining check_repo item
for esmfold2 is the model-doc page, which is a separate pending task.
The tiny test config encodes two real sizing constraints discovered via the
forward smoke: 3D RoPE needs 3*n_spatial + n_uid <= head_dim//2, and
inputs.d_inputs == 67 + d_token//2 == diffusion_module.c_s_inputs.
7 non-slow tests pass; ruff + check_config_attributes + check_docstrings clean.
Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
