Skip to content

Adjust InjectorSurrogate module based on LUMETorchModel updates#20

Open
pluflou wants to merge 27 commits intoslaclab:injector-surrogatefrom
pluflou:update-torch
Open

Adjust InjectorSurrogate module based on LUMETorchModel updates#20
pluflou wants to merge 27 commits intoslaclab:injector-surrogatefrom
pluflou:update-torch

Conversation

@pluflou
Copy link
Copy Markdown
Contributor

@pluflou pluflou commented Mar 30, 2026

This pull request refactors the InjectorSurrogate class to align with changes in LUMETorchModel.

  • Refactored the internal state management by introducing the _cache dictionary (as in LUMETorchModel), removing the old _state and update_state methods. The model is initialized with default values (empty set call).
  • The model now avoids mutating shared internal state unsafely (self._cache = self.model._cache was aliasing a private attribute).
  • Updated the supported_variables property to avoid mutating the wrapped model's metadata and to use a more explicit interface.
  • Added helper functions _tensor_to_numpy and _to_python_scalar to safely convert tensors to NumPy arrays and Python scalars.
  • Updated examples/injector_surrogate.ipynb to use the new output_beam key instead of beam_output.
  • Linting changes to many files.
  • Fixed import issue for the YAML model config that's stored in the subtrees. It now temporarily adjusts the paths of the resources at instantiation then deletes the modified YAML. This allows us to keep the surrogate model configs up to date with the upstream repos without having to rewrite the paths.
    • Deleted the old .submodules dir because the naming is confusing (they are not submodules)
    • Added a subtrees/repo_name directory that clarifies what is in there
    • Adjusted model to find the files based on where it's running (CI vs local)

@pluflou pluflou marked this pull request as draft March 31, 2026 00:21
@pluflou pluflou marked this pull request as ready for review March 31, 2026 19:00
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.

2 participants