[Enhancement] enable array API support in EmpiricialCovariance and IncrementalEmpiricalCovariance (#2207)

* fix onedal side

* add SPMD interface

* add new assert_all_finite

* change incremental algo

* remove old code

* readd deletion

* missing?

* Update test_covariance.py

* fix error in spmd interface

* fixes

* remove changes

* fix

* fix preview

* fixes

* forgotten import

* remove import

* fix squeeze

* Update test_covariance.py

* Update covariance.py

* Update incremental_covariance.py

* Update incremental_covariance.py

* try to fix issues

* updates to try and solve score

* fixes for mahalanobis

* Update incremental_covariance.py

* Update incremental_covariance.py

* Update incremental_covariance.py

* Update incremental_covariance.py

* Update incremental_covariance.py

* Update incremental_covariance.py

* expand array API support in sklearnex

* fix spelling mistake

* remove print statement

* fix issues with respect to dpctl dpnp

* Update covariance.py

* fix error_norm and score

* try to fix mahal

* fixes

* fixes

* fixes again

* try again

* fixes for test case failures

* try again

* try to remove regression

* fix mistakes

* updates

* fix covariance

* standardize validate_params

* Update covariance.py

* Update covariance.py

* fix issues with double import

* Update covariance.py

* Update covariance.py

* Update _array_api.py

* Update _array_api.py

* Update _array_api.py

* Update _array_api.py

* Update incremental_covariance.py

* Update incremental_covariance.py

* Update covariance.py

* Update incremental_covariance.py

* Update covariance.py

* Update covariance.py

* Update incremental_covariance.py

* Update incremental_covariance.py

* Update incremental_covariance.py

* Update incremental_covariance.py

* Update covariance.py

* Update _array_api.py

* Update covariance.py

* Update incremental_covariance.py

* Update test_incremental_covariance_spmd.py

* Update covariance.py

* Update test_incremental_covariance_spmd.py

* Update covariance.py

* Update incremental_covariance.py

* Update covariance.py

* Update test_covariance_spmd.py

* Update test_covariance_spmd.py

* Update test_covariance_spmd.py

* Update test_covariance_spmd.py

* Update test_covariance_spmd.py

* many requisite changes

* Update test_covariance.py

* force output to a 0d

* try to fix test'

* updates

Author: icfaust

onedal/covariance/covariance.py

@@ -20,7 +20,6 @@
 from daal4py.sklearn._utils import daal_check_version
 from onedal._device_offload import supports_queue
 from onedal.common._backend import bind_default_backend
-from onedal.utils.validation import _check_array
 
 from .._config import _get_config
 from ..common.hyperparameters import get_hyperparameters
@@ -101,13 +100,7 @@ def fit(self, X, y=None, queue=None):
         self : object
             Returns the instance itself.
         """
-        use_raw_input = _get_config()["use_raw_input"] is True
-        sua_iface, xp, _ = _get_sycl_namespace(X)
-        if use_raw_input and sua_iface:
-            queue = X.sycl_queue
 
-        if not use_raw_input:
-            X = _check_array(X, dtype=[np.float64, np.float32])
         X_table = to_table(X, queue=queue)
 
         params = self._get_onedal_params(X_table.dtype)
@@ -123,6 +116,6 @@ def fit(self, X, y=None, queue=None):
                 from_table(result.cov_matrix, like=X) * (X.shape[0] - 1) / X.shape[0]
             )
 
-        self.location_ = xp.squeeze(from_table(result.means, like=X))
+        self.location_ = from_table(result.means, like=X)[0, ...]
 
         return self

onedal/covariance/incremental_covariance.py

@@ -17,12 +17,12 @@
 import numpy as np
 
 from daal4py.sklearn._utils import daal_check_version
-from onedal._device_offload import supports_queue
-from onedal.common._backend import bind_default_backend
-from onedal.utils import _sycl_queue_manager as QM
 
 from .._config import _get_config
-from ..datatypes import from_table, to_table
+from .._device_offload import supports_queue
+from ..common._backend import bind_default_backend
+from ..datatypes import from_table, return_type_constructor, to_table
+from ..utils import _sycl_queue_manager as QM
 from ..utils._array_api import _get_sycl_namespace
 from ..utils.validation import _check_array
 from .covariance import BaseEmpiricalCovariance
@@ -74,6 +74,7 @@ def finalize_compute(self, params, partial_result): ...
     def _reset(self):
         self._need_to_finalize = False
         self._queue = None
+        self._outtype = None
         self._partial_result = self.partial_compute_result()
 
     def __getstate__(self):
@@ -108,15 +109,10 @@ def partial_fit(self, X, y=None, queue=None):
         self : object
             Returns the instance itself.
         """
-        use_raw_input = _get_config()["use_raw_input"] is True
-        sua_iface, _, _ = _get_sycl_namespace(X)
-
-        if use_raw_input and sua_iface:
-            queue = X.sycl_queue
-        if not use_raw_input:
-            X = _check_array(X, dtype=[np.float64, np.float32], ensure_2d=True)
 
         self._queue = queue
+        if not self._outtype:
+            self._outtype = return_type_constructor(X)
         X_table = to_table(X, queue=queue)
 
         if not hasattr(self, "_dtype"):
@@ -125,8 +121,6 @@ def partial_fit(self, X, y=None, queue=None):
         params = self._get_onedal_params(self._dtype)
         self._partial_result = self.partial_compute(params, self._partial_result, X_table)
         self._need_to_finalize = True
-        # store the queue for when we finalize
-        self._queue = queue
 
     def finalize_fit(self):
         """Finalize covariance matrix from the current `_partial_result`.
@@ -143,13 +137,14 @@ def finalize_fit(self):
             with QM.manage_global_queue(self._queue):
                 result = self.finalize_compute(params, self._partial_result)
 
-            if daal_check_version((2024, "P", 1)) or (not self.bias):
-                self.covariance_ = from_table(result.cov_matrix)
-            else:
+            self.covariance_ = from_table(result.cov_matrix, like=self._outtype)
+
+            if self.bias and not daal_check_version((2024, "P", 1)):
                 n_rows = self._partial_result.partial_n_rows
-                self.covariance_ = from_table(result.cov_matrix) * (n_rows - 1) / n_rows
+                self.covariance_ *= (n_rows - 1) / n_rows
 
-            self.location_ = from_table(result.means).ravel()
+            self.location_ = from_table(result.means, like=self._outtype)[0, ...]
+            self._outtype = None
 
             self._need_to_finalize = False
 

onedal/covariance/tests/test_covariance.py

@@ -25,23 +25,23 @@
 def test_onedal_import_covariance(queue):
     from onedal.covariance import EmpiricalCovariance
 
-    X = np.array([[0, 1], [0, 1]])
+    X = np.array([[0, 1], [0, 1]], dtype=np.float64)
     result = EmpiricalCovariance().fit(X, queue=queue)
     expected_covariance = np.array([[0, 0], [0, 0]])
     expected_means = np.array([0, 1])
 
     assert_allclose(expected_covariance, result.covariance_)
     assert_allclose(expected_means, result.location_)
 
-    X = np.array([[1, 2], [3, 6]])
+    X = np.array([[1, 2], [3, 6]], dtype=np.float64)
     result = EmpiricalCovariance().fit(X, queue=queue)
     expected_covariance = np.array([[2, 4], [4, 8]])
     expected_means = np.array([2, 4])
 
     assert_allclose(expected_covariance, result.covariance_)
     assert_allclose(expected_means, result.location_)
 
-    X = np.array([[1, 2], [3, 6]])
+    X = np.array([[1, 2], [3, 6]], dtype=np.float64)
     result = EmpiricalCovariance(bias=True).fit(X, queue=queue)
     expected_covariance = np.array([[1, 2], [2, 4]])
     expected_means = np.array([2, 4])