Use array_api in Gaussian Mixture Models

sklearn/_config.py

@@ -9,6 +9,7 @@
     "working_memory": int(os.environ.get("SKLEARN_WORKING_MEMORY", 1024)),
     "print_changed_only": True,
     "display": "text",
+    "array_api_dispatch": False,
 }
 _threadlocal = threading.local()
 
@@ -40,7 +41,11 @@ def get_config():
 
 
 def set_config(
-    assume_finite=None, working_memory=None, print_changed_only=None, display=None
+    assume_finite=None,
+    working_memory=None,
+    print_changed_only=None,
+    display=None,
+    array_api_dispatch=None,
 ):
     """Set global scikit-learn configuration
 
@@ -95,11 +100,18 @@ def set_config(
         local_config["print_changed_only"] = print_changed_only
     if display is not None:
         local_config["display"] = display
+    if array_api_dispatch is not None:
+        local_config["array_api_dispatch"] = array_api_dispatch
 
 
 @contextmanager
 def config_context(
-    *, assume_finite=None, working_memory=None, print_changed_only=None, display=None
+    *,
+    assume_finite=None,
+    working_memory=None,
+    print_changed_only=None,
+    display=None,
+    array_api_dispatch=None,
 ):
     """Context manager for global scikit-learn configuration.
 
@@ -171,6 +183,7 @@ def config_context(
         working_memory=working_memory,
         print_changed_only=print_changed_only,
         display=display,
+        array_api_dispatch=array_api_dispatch,
     )
 
     try:

sklearn/mixture/_base.py

@@ -9,13 +9,13 @@
 from time import time
 
 import numpy as np
-from scipy.special import logsumexp
 
 from .. import cluster
 from ..base import BaseEstimator
 from ..base import DensityMixin
 from ..exceptions import ConvergenceWarning
 from ..utils import check_random_state
+from ..utils._array_api import get_namespace, logsumexp
 from ..utils.validation import check_is_fitted
 
 
@@ -136,6 +136,7 @@ def _initialize_parameters(self, X, random_state):
             used for the method chosen to initialize the parameters.
         """
         n_samples, _ = X.shape
+        np, _ = get_namespace(X)
 
         if self.init_params == "kmeans":
             resp = np.zeros((n_samples, self.n_components))
@@ -149,7 +150,8 @@ def _initialize_parameters(self, X, random_state):
             resp[np.arange(n_samples), label] = 1
         elif self.init_params == "random":
             resp = random_state.rand(n_samples, self.n_components)
-            resp /= resp.sum(axis=1)[:, np.newaxis]
+            resp = np.asarray(resp)
+            resp /= np.reshape(np.sum(resp, axis=1), (-1, 1))
         else:
             raise ValueError(
                 "Unimplemented initialization method '%s'" % self.init_params
@@ -225,6 +227,7 @@ def fit_predict(self, X, y=None):
         labels : array, shape (n_samples,)
             Component labels.
         """
+        np, _ = get_namespace(X)
         X = self._validate_data(X, dtype=[np.float64, np.float32], ensure_min_samples=2)
         if X.shape[0] < self.n_components:
             raise ValueError(
@@ -291,7 +294,7 @@ def fit_predict(self, X, y=None):
         # for any value of max_iter and tol (and any random_state).
         _, log_resp = self._e_step(X)
 
-        return log_resp.argmax(axis=1)
+        return np.argmax(log_resp, axis=1)
 
     def _e_step(self, X):
         """E step.
@@ -309,6 +312,7 @@ def _e_step(self, X):
             Logarithm of the posterior probabilities (or responsibilities) of
             the point of each sample in X.
         """
+        np, _ = get_namespace(X)
         log_prob_norm, log_resp = self._estimate_log_prob_resp(X)
         return np.mean(log_prob_norm), log_resp
 
@@ -527,11 +531,15 @@ def _estimate_log_prob_resp(self, X):
         log_responsibilities : array, shape (n_samples, n_components)
             logarithm of the responsibilities
         """
+        np, is_array_api = get_namespace(X)
         weighted_log_prob = self._estimate_weighted_log_prob(X)
         log_prob_norm = logsumexp(weighted_log_prob, axis=1)
-        with np.errstate(under="ignore"):
-            # ignore underflow
-            log_resp = weighted_log_prob - log_prob_norm[:, np.newaxis]
+        if is_array_api:
+            log_resp = weighted_log_prob - np.reshape(log_prob_norm, (-1, 1))
+        else:
+            with np.errstate(under="ignore"):
+                # ignore underflow
+                log_resp = weighted_log_prob - log_prob_norm[:, np.newaxis]
         return log_prob_norm, log_resp
 
     def _print_verbose_msg_init_beg(self, n_init):

sklearn/mixture/_gaussian_mixture.py

@@ -5,12 +5,16 @@
 # License: BSD 3 clause
 
 import numpy as np
+from math import log
+from functools import partial
 
 from scipy import linalg
+import scipy
 
 from ._base import BaseMixture, _check_shape
 from ..utils import check_array
 from ..utils.extmath import row_norms
+from ..utils._array_api import get_namespace
 
 
 ###############################################################################
@@ -171,12 +175,13 @@ def _estimate_gaussian_covariances_full(resp, X, nk, means, reg_covar):
     covariances : array, shape (n_components, n_features, n_features)
         The covariance matrix of the current components.
     """
+    np, _ = get_namespace(resp, X, nk)
     n_components, n_features = means.shape
     covariances = np.empty((n_components, n_features, n_features))
     for k in range(n_components):
-        diff = X - means[k]
-        covariances[k] = np.dot(resp[:, k] * diff.T, diff) / nk[k]
-        covariances[k].flat[:: n_features + 1] += reg_covar
+        diff = X - means[k, :]
+        covariances[k, :, :] = ((resp[:, k] * diff.T) @ diff) / nk[k]
+        np.reshape(covariances[k, :, :], (-1,))[:: n_features + 1] += reg_covar
     return covariances
 
 
@@ -286,8 +291,9 @@ def _estimate_gaussian_parameters(X, resp, reg_covar, covariance_type):
         The covariance matrix of the current components.
         The shape depends of the covariance_type.
     """
-    nk = resp.sum(axis=0) + 10 * np.finfo(resp.dtype).eps
-    means = np.dot(resp.T, X) / nk[:, np.newaxis]
+    np, _ = get_namespace(X, resp)
+    nk = np.sum(resp, axis=0) + 10 * np.finfo(resp.dtype).eps
+    means = resp.T @ X / np.reshape(nk, (-1, 1))
     covariances = {
         "full": _estimate_gaussian_covariances_full,
         "tied": _estimate_gaussian_covariances_tied,
@@ -321,27 +327,31 @@ def _compute_precision_cholesky(covariances, covariance_type):
         "or collapsed samples). Try to decrease the number of components, "
         "or increase reg_covar."
     )
+    np, is_array_api = get_namespace(covariances)
+    if is_array_api:
+        cholesky = np.linalg.cholesky
+        solve = np.linalg.solve
+    else:
+        cholesky = partial(scipy.linalg.cholesky, lower=True)
+        solve = partial(scipy.linalg.solve_triangular, lower=True)
 
     if covariance_type == "full":
         n_components, n_features, _ = covariances.shape
         precisions_chol = np.empty((n_components, n_features, n_features))
-        for k, covariance in enumerate(covariances):
+        for k in range(n_components):
             try:
-                cov_chol = linalg.cholesky(covariance, lower=True)
+                cov_chol = cholesky(covariances[k, :, :])
             except linalg.LinAlgError:
                 raise ValueError(estimate_precision_error_message)
-            precisions_chol[k] = linalg.solve_triangular(
-                cov_chol, np.eye(n_features), lower=True
-            ).T
+            precisions_chol[k, :, :] = solve(cov_chol, np.eye(n_features)).T
+
     elif covariance_type == "tied":
         _, n_features = covariances.shape
         try:
-            cov_chol = linalg.cholesky(covariances, lower=True)
+            cov_chol = cholesky(covariances)
         except linalg.LinAlgError:
             raise ValueError(estimate_precision_error_message)
-        precisions_chol = linalg.solve_triangular(
-            cov_chol, np.eye(n_features), lower=True
-        ).T
+        precisions_chol = linalg.solve(cov_chol, np.eye(n_features)).T
     else:
         if np.any(np.less_equal(covariances, 0.0)):
             raise ValueError(estimate_precision_error_message)
@@ -373,11 +383,11 @@ def _compute_log_det_cholesky(matrix_chol, covariance_type, n_features):
     log_det_precision_chol : array-like of shape (n_components,)
         The determinant of the precision matrix for each component.
     """
+    np, _ = get_namespace(matrix_chol)
     if covariance_type == "full":
         n_components, _, _ = matrix_chol.shape
-        log_det_chol = np.sum(
-            np.log(matrix_chol.reshape(n_components, -1)[:, :: n_features + 1]), 1
-        )
+        matrix_col_reshape = np.reshape(matrix_chol, (n_components, -1))
+        log_det_chol = np.sum(np.log(matrix_col_reshape[:, :: n_features + 1]), axis=1)
 
     elif covariance_type == "tied":
         log_det_chol = np.sum(np.log(np.diag(matrix_chol)))
@@ -413,6 +423,7 @@ def _estimate_log_gaussian_prob(X, means, precisions_chol, covariance_type):
     -------
     log_prob : array, shape (n_samples, n_components)
     """
+    np, _ = get_namespace(X, means, precisions_chol)
     n_samples, n_features = X.shape
     n_components, _ = means.shape
     # The determinant of the precision matrix from the Cholesky decomposition
@@ -423,8 +434,10 @@ def _estimate_log_gaussian_prob(X, means, precisions_chol, covariance_type):
 
     if covariance_type == "full":
         log_prob = np.empty((n_samples, n_components))
-        for k, (mu, prec_chol) in enumerate(zip(means, precisions_chol)):
-            y = np.dot(X, prec_chol) - np.dot(mu, prec_chol)
+        for k in range(n_components):
+            mu = means[k, :]
+            prec_chol = precisions_chol[k, :, :]
+            y = X @ prec_chol - mu @ prec_chol
             log_prob[:, k] = np.sum(np.square(y), axis=1)
 
     elif covariance_type == "tied":
@@ -450,7 +463,7 @@ def _estimate_log_gaussian_prob(X, means, precisions_chol, covariance_type):
         )
     # Since we are using the precision of the Cholesky decomposition,
     # `- 0.5 * log_det_precision` becomes `+ log_det_precision_chol`
-    return -0.5 * (n_features * np.log(2 * np.pi) + log_prob) + log_det
+    return -0.5 * (n_features * log(2 * np.pi) + log_prob) + log_det
 
 
 class GaussianMixture(BaseMixture):
@@ -742,6 +755,7 @@ def _m_step(self, X, log_resp):
             the point of each sample in X.
         """
         n_samples, _ = X.shape
+        np, _ = get_namespace(X, log_resp)
         self.weights_, self.means_, self.covariances_ = _estimate_gaussian_parameters(
             X, np.exp(log_resp), self.reg_covar, self.covariance_type
         )
@@ -756,6 +770,7 @@ def _estimate_log_prob(self, X):
         )
 
     def _estimate_log_weights(self):
+        np, _ = get_namespace(self.weights_)
         return np.log(self.weights_)
 
     def _compute_lower_bound(self, _, log_prob_norm):

sklearn/utils/_array_api.py

@@ -0,0 +1,79 @@
+"""Tools to support array_api."""
+import numpy as np
+from scipy.special import logsumexp as sp_logsumexp
+from .._config import get_config
+
+
+def get_namespace(*xs):
+    # `xs` contains one or more arrays, or possibly Python scalars (accepting
+    # those is a matter of taste, but doesn't seem unreasonable).
+    # Returns a tuple: (array_namespace, is_array_api)
+
+    if not get_config()["array_api_dispatch"]:
+        return np, False
+
+    namespaces = {
+        x.__array_namespace__() if hasattr(x, "__array_namespace__") else None
+        for x in xs
+        if not isinstance(x, (bool, int, float, complex))
+    }
+
+    if not namespaces:
+        # one could special-case np.ndarray above or use np.asarray here if
+        # older numpy versions need to be supported.
+        raise ValueError("Unrecognized array input")
+
+    if len(namespaces) != 1:
+        raise ValueError(f"Multiple namespaces for array inputs: {namespaces}")
+
+    (xp,) = namespaces
+    if xp is None:
+        # Use numpy as default
+        return np, False
+
+    return xp, True
+
+
+def logsumexp(a, axis=None, b=None, keepdims=False, return_sign=False):
+    np, is_array_api = get_namespace(a)
+
+    # Use SciPy if a is an ndarray
+    if not is_array_api:
+        return sp_logsumexp(
+            a, axis=axis, b=b, keepdims=keepdims, return_sign=return_sign
+        )
+
+    if b is not None:
+        a, b = np.broadcast_arrays(a, b)
+        if np.any(b == 0):
+            a = a + 0.0  # promote to at least float
+            a[b == 0] = -np.inf
+
+    a_max = np.max(a, axis=axis, keepdims=True)
+
+    if a_max.ndim > 0:
+        a_max[~np.isfinite(a_max)] = 0
+    elif not np.isfinite(a_max):
+        a_max = 0
+
+    if b is not None:
+        b = np.asarray(b)
+        tmp = b * np.exp(a - a_max)
+    else:
+        tmp = np.exp(a - a_max)
+
+    # suppress warnings about log of zero
+    s = np.sum(tmp, axis=axis, keepdims=keepdims)
+    if return_sign:
+        sgn = np.sign(s)
+        s *= sgn  # /= makes more sense but we need zero -> zero
+    out = np.log(s)
+
+    if not keepdims:
+        a_max = np.squeeze(a_max, axis=axis)
+    out += a_max
+
+    if return_sign:
+        return out, sgn
+    else:
+        return out