use GPUs for GLMs Pt.2

Author: thierrymoudiki

examples/glm_regression_quantile_gpopt.py

@@ -146,29 +146,13 @@ def compute_XB(self, X, beta=None, row_index=None):
         return mo.safe_sparse_dot(X[row_index, :], self.beta_, backend=self.backend)
 
     def compute_XB2(self, X, beta=None, row_index=None):
-        if self.backend != "cpu":
-            raise NotImplementedError(
-                "GLM.compute_XB is only implemented for backend='cpu'"
-            )
         def f00(X):
-            if self.backend != "cpu":
-                raise NotImplementedError(
-                "GLM.compute_XB is only implemented for backend='cpu'"
-            )
             return mo.safe_sparse_dot(X, self.beta_, backend=self.backend)
 
         def f01(X):
-            if self.backend != "cpu":
-                raise NotImplementedError(
-                "GLM.compute_XB is only implemented for backend='cpu'"
-            )
             return mo.safe_sparse_dot(X[row_index, :], self.beta_, backend=self.backend)
 
         def f11(X):
-            if self.backend != "cpu":
-                raise NotImplementedError(
-                "GLM.compute_XB is only implemented for backend='cpu'"
-            )
             return mo.safe_sparse_dot(X[row_index, :], beta, backend=self.backend)
 
         def f10(X):
@@ -186,10 +170,6 @@ def f10(X):
         return h_result[result_code](X)
 
     def penalty(self, beta1, beta2, lambda1, lambda2, alpha1, alpha2):
-        if self.backend != "cpu":
-            raise NotImplementedError(
-                "GLM.compute_XB is only implemented for backend='cpu'"
-            )
         res = lambda1 * (
             0.5 * (1 - alpha1) * np.sum(np.square(beta1))
             + alpha1 * np.sum(np.abs(beta1))
@@ -201,10 +181,6 @@ def penalty(self, beta1, beta2, lambda1, lambda2, alpha1, alpha2):
         return res
 
     def compute_penalty(self, group_index, beta):
-        if self.backend != "cpu":
-            raise NotImplementedError(
-                "GLM.compute_XB is only implemented for backend='cpu'"
-            )
         return self.penalty(
             beta1=beta[0:group_index],
             beta2=beta[group_index: len(beta)],

nnetsauce/glm/glm.py

@@ -10,7 +10,10 @@
 from sklearn.base import ClassifierMixin
 from ..optimizers import Optimizer
 from scipy.special import logsumexp, expit, erf
-
+try: 
+    import jax.numpy as jnp
+except ImportError:
+    pass
 
 class GLMClassifier(GLM, ClassifierMixin):
     """Generalized 'linear' models using quasi-randomized networks (classification)
@@ -244,7 +247,10 @@ def fit(self, X, y, **kwargs):
         Y = self.optimizer.one_hot_encode(output_y, self.n_classes)
 
         # initialization
-        beta_ = np.linalg.lstsq(scaled_Z, Y, rcond=None)[0]
+        if self.backend == "cpu":
+            beta_ = np.linalg.lstsq(scaled_Z, Y, rcond=None)[0]
+        else:
+            beta_ = jnp.linalg.lstsq(scaled_Z, Y, rcond=None)[0]
 
         # optimization
         # fit(self, loss_func, response, x0, **kwargs):

nnetsauce/glm/glmClassifier.py

@@ -9,6 +9,10 @@
 from ..utils import matrixops as mo
 from sklearn.base import RegressorMixin
 from ..optimizers import Optimizer
+try: 
+    import jax.numpy as jnp
+except ImportError:
+    pass
 
 
 class GLMRegressor(GLM, RegressorMixin):
@@ -239,7 +243,10 @@ def fit(self, X, y, **kwargs):
 
         centered_y, scaled_Z = self.cook_training_set(y=y, X=X, **kwargs)
         # initialization
-        beta_ = np.linalg.lstsq(scaled_Z, centered_y, rcond=None)[0]
+        if self.backend == "cpu":
+            beta_ = np.linalg.lstsq(scaled_Z, centered_y, rcond=None)[0]
+        else:
+            beta_ = jnp.linalg.lstsq(scaled_Z, centered_y, rcond=None)[0]
         # optimization
         # fit(self, loss_func, response, x0, **kwargs):
         # loss_func(self, beta, group_index, X, y,