Skip to content

Remove legacy LinearReferences class and update to use ElementReferences#1796

Open
lbluque wants to merge 4 commits intomainfrom
element-refs-cleanup
Open

Remove legacy LinearReferences class and update to use ElementReferences#1796
lbluque wants to merge 4 commits intomainfrom
element-refs-cleanup

Conversation

@lbluque
Copy link
Contributor

@lbluque lbluque commented Feb 18, 2026

Summary

The LinearReferences and ElementReferences classes were essentially duplicated functionality. This removes the LinearReferences class and keeps ElementReferences - which is being used in UMA models.

  • Remove the LinearReferences class from element_references.py, consolidating all element reference functionality into the single ElementReferences class already used by the training pipeline and YAML configs.
  • Update fit_linear_references, create_element_references, and load_references_from_config to return ElementReferences instead of LinearReferences.
  • Update normalizer.py to use the ElementReferences API (apply_refs) instead of the removed LinearReferences.dereference method.
  • Add a getattr fallback in ElementReferences.compute_references so it works both after a model forward pass (where batch_full/atomic_numbers_full are set by the backbone) and in non-model contexts like reference fitting and normalizer estimation.
  • Re-enable and fix the previously skipped tests in test_element_references.py to use the ElementReferences API.

@meta-cla meta-cla bot added the cla signed label Feb 18, 2026
@lbluque lbluque added deprecation Code deprecation patch Patch version release and removed cla signed labels Feb 18, 2026
@lbluque lbluque requested review from rayg1234 and wood-b February 18, 2026 21:08
@meta-cla meta-cla bot added the cla signed label Feb 18, 2026
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

cla signed deprecation Code deprecation patch Patch version release

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant