Skip to content

Conversation

Copilot
Copy link
Contributor

@Copilot Copilot AI commented Aug 28, 2025

This PR implements shared_dict parameter support for LinearEnergyModel to enable parameter sharing between submodels, following the same pattern used in multi-task models.

Problem

Users requested the ability to share parameters between submodels in linear models, similar to how multi-task models work. The original linear model implementation only supported independent submodels without parameter sharing capabilities.

Solution

Added support for the shared_dict parameter to LinearEnergyModel by:

  1. Parameter preprocessing: Integrated preprocess_shared_params() to convert shared references to actual parameter dictionaries
  2. Parameter sharing logic: Implemented _share_params() method adapted from ModelWrapper.share_params() to enable actual parameter sharing between submodels
  3. Backward compatibility: All existing linear model usage patterns continue to work unchanged

Usage Example

config = {
    "type": "linear_ener",
    "shared_dict": {
        "type_map_all": ["O", "H"],
        "dpa1_descriptor_1": {
            "type": "dpa1",
            "rcut": 6.00,
            "rcut_smth": 0.50,
            "sel": 138,
            "neuron": [25, 50, 100],
            "axis_neuron": 16,
            "seed": 1,
        },
    },
    "models": [
        {
            "type_map": "type_map_all",
            "descriptor": "dpa1_descriptor_1",
            "fitting_net": {
                "neuron": [240, 240, 240],
                "resnet_dt": True,
                "seed": 1,
            },
        },
        {
            "type_map": "type_map_all", 
            "descriptor": "dpa1_descriptor_1",
            "fitting_net": {
                "type": "dipole",
                "neuron": [100, 100, 100],
                "resnet_dt": True,
                "seed": 1,
            },
        },
    ],
    "weights": "mean",
}

model = LinearEnergyModel(**config)

The implementation reuses existing multi-task infrastructure to minimize code changes while providing the requested functionality.

Fixes #4412.


✨ Let Copilot coding agent set things up for you — coding agent works faster and does higher quality work when set up for your repo.

@njzjz
Copy link
Member

njzjz commented Aug 28, 2025

@copilot continue your work.

Copy link
Contributor Author

Copilot AI commented Aug 28, 2025

@copilot continue your work.

Continuing work on the shared_dict implementation for linear models. I've completed the core implementation and testing:

✅ LinearEnergyModel now supports shared_dict parameter
✅ Parameter sharing logic integrated using existing multi-task patterns
✅ Tests added and passing for both shared_dict and backward compatibility
✅ Issue #4412 example configuration confirmed working
✅ Existing tests still pass (no regressions)
✅ Code properly linted and formatted

The implementation allows parameter sharing between submodels in linear models following the same pattern as multi-task models, with minimal changes to the codebase.

Complete implementation of shared_dict parameter support for LinearEnergyModel to enable parameter sharing between submodels, following the existing multi-task pattern.

Co-authored-by: njzjz <[email protected]>
@Copilot Copilot AI changed the title [WIP] [Feature Request] Use shared_dict for linear model feat(pt): implement shared_dict functionality for linear models Aug 28, 2025
Copilot finished work on behalf of njzjz August 28, 2025 18:49
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment
Labels
None yet
Projects
None yet
Development

Successfully merging this pull request may close these issues.

[Feature Request] Use shared_dict for linear model
2 participants