drafting export_torchscript

Author: agoscinski

src/equisolve/nn/module_tensor.py

@@ -242,6 +242,36 @@ def from_module(
         module = torch.nn.Linear(in_features, out_features, bias, device, dtype)
         return ModuleTensorMap.from_module(in_keys, module, many_to_one, out_tensor)
 
+    @classmethod
+    def from_weights(
+        cls,
+        weights: TensorMap,
+        bias: Optional[TensorMap] = None
+    ):
+        """
+        :param weights:
+            The weight tensor map from which we create the linear modules.
+
+        :param bias:
+            The weight tensor map from which we create the linear layers.
+        """
+        module_map = ModuleDict()
+        for key, weights_block in weights.items():
+            module_key = ModuleTensorMap.module_key(key)
+            module = torch.nn.Linear(
+                len(weights_block.samples),
+                len(weights_block.properties),
+                bias=False,
+                device=weights_block.values.device,
+                dtype=weights_block.values.dtype,
+            )
+            module.weight = torch.nn.Parameter(weights_block.values)
+            if bias is not None:
+                module.bias = torch.nn.Parameter(bias.block(key).values)
+            module_map[module_key] = module
+
+        return ModuleTensorMap(module_map, weights)
+
     def forward(self, tensor: TensorMap) -> TensorMap:
         # added to appear in doc, :inherited-members: is not compatible with torch
         return super().forward(tensor)

src/equisolve/numpy/models/linear_model.py

@@ -13,13 +13,12 @@
 import scipy.linalg
 from metatensor import Labels, TensorBlock, TensorMap
 
-from ... import HAS_TORCH
 from ...module import NumpyModule, _Estimator
 from ...utils.metrics import rmse
 from ..utils import array_from_block
 
 
-class _Ridge(_Estimator):
+class Ridge(_Estimator):
     r"""Linear least squares with l2 regularization for :class:`metatensor.Tensormap`'s.
 
     Weights :math:`w` are calculated according to
@@ -352,21 +351,15 @@ def score(self, X: TensorMap, y: TensorMap, parameter_key: str) -> float:
         y_pred = self.predict(X)
         return rmse(y, y_pred, parameter_key)
 
-
-class NumpyRidge(_Ridge, NumpyModule):
-    def __init__(self) -> None:
-        NumpyModule.__init__(self)
-        _Ridge.__init__(self)
-
-
-if HAS_TORCH:
-    import torch
-
-    class TorchRidge(_Ridge, torch.nn.Module):
-        def __init__(self) -> None:
-            torch.nn.Module.__init__(self)
-            _Ridge.__init__(self)
-
-    Ridge = TorchRidge
-else:
-    Ridge = NumpyRidge
+    def export_torchscript(self):
+        from ... import HAS_METATENSOR_TORCH
+        if not HAS_METATENSOR_TORCH:
+            raise ImportError(
+                "To export your model to TorchScript torch needs to be installed. "
+                "Please install torch, then reimport equisolve or "
+                "use equisolve.refresh_global_flags()."
+            )
+        from ...nn import Linear
+        # TODO should be metatensor.torch
+        weights = metatensor.to(self.weights, backend="torch")
+        return Linear.from_weights(weights)

tests/equisolve_tests/numpy/models/linear_model.py

@@ -13,6 +13,8 @@
 
 from equisolve.numpy.models import Ridge
 
+from equisolve import HAS_METATENSOR_TORCH
+
 from ..utilities import tensor_to_tensormap
 
 
@@ -79,6 +81,29 @@ def equisolve_solver_from_numpy_arrays(
         clf.fit(X=X, y=y, alpha=alpha, sample_weight=sw, solver=solver)
         return clf
 
+    @pytest.mark.skipif(
+        not (HAS_METATENSOR_TORCH), reason="requires metatensor-torch to be run"
+    )
+    def test_export_torchscript(self):
+        """Test if ridge is working and all shapes are converted correctly.
+        Test is performed for two blocks.
+        """
+
+        num_targets = 50
+        num_properties = 5
+
+        # Create input values
+        X_arr = self.rng.random([2, num_targets, num_properties])
+        y_arr = self.rng.random([2, num_targets, 1])
+
+        X = tensor_to_tensormap(X_arr)
+        y = tensor_to_tensormap(y_arr)
+
+        clf = Ridge()
+        clf.fit(X=X, y=y)
+        module = clf.export_torchscript()
+        # TODO test torchscriptability and check if output match with original Ridge
+
     num_properties = np.array([91])
     num_targets = np.array([1000])
     means = np.array([-0.5, 0, 0.1])