TorchScript compatibility

[agoscinski]

Discussion has been moved to issue #34

A draft PR to figure out what needs to be actually done on the equisolve side to add TorchScript support. This will be probably splitted int smaller sub PRs

Replacing numpy operations with equistore operations in our transformers/estimators

• Ridge (for predict Improvements to Ridge #11 already done)
• StandardScaler (issue Rewrite StandardScaler in terms of equistore operations #18)

I think it would be easier for now if we limit ourselves to the forward functions, since in the fit function operations happen that are cannot be completely reproduced with equistore operations and I am not sure how to split them into equistore.operations.

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

Removing numpy torch layer

• small refactor removing the top level with the numpy and torch folder, because they get merged by using equistore.operations

---

📚 Documentation preview 📚: https://equisolve--24.org.readthedocs.build/en/24/