Skip to content
This repository has been archived by the owner on Apr 24, 2024. It is now read-only.

TorchScript compatibility equisolve side #34

Open
6 tasks
agoscinski opened this issue Mar 6, 2023 · 0 comments
Open
6 tasks

TorchScript compatibility equisolve side #34

agoscinski opened this issue Mar 6, 2023 · 0 comments

Comments

@agoscinski
Copy link
Collaborator

agoscinski commented Mar 6, 2023

Moved this from initial PR #24

This is basically issue metatensor/metatensor#133 for the equisolve side

Replacing numpy operations with equistore operations in our transformers/estimators

I think it would be easier for now if we limit ourselves to the forward functions, since in the fit function operations happen that cannot be completely reproduced with equistore operations at the moment. But we keep the long term plan to replace it

Add optional torch.nn.Module inheritance on the transformers and estimators

Add this inheritance class if torch is installed. That way we can get a class that is torch.jit.compile'able in case TensorMaps with torch.tensor's are used. In the cases where torch is not installed or a TensorMap with numpy.ndarray's are used, then at least it can be dumped as pickle object.

  • We need to check if the torch.jit.compile store gives a good error message when the object contains torch.tensor objects (e.g. weight matrix in Ridge) to store data, so the user knows why torch.jit.compile does not work in these cases
  • Because torch does not support all options that scipy.linalg.solve does, we might want to do a fitting on TensorMaps with numpy.ndarray's and then change the data to torch.tensor's to later use the estimator/transformer as a TorchScript -- After discussion mentioned in Change linear system solver from numpy to scipy metatensor/metatensor#152 we might be able to recreate the other solvers with torch operations. The numerical stability might differ between numpy and torch, but at least the same methods should be available

Removing numpy torch layer

  • small refactor removing the top level with the numpy and torch folder, because they get merged by using equistore.operations to modules.
@agoscinski agoscinski mentioned this issue Mar 6, 2023
5 tasks
Sign up for free to subscribe to this conversation on GitHub. Already have an account? Sign in.
Labels
None yet
Projects
None yet
Development

No branches or pull requests

1 participant