fix: fallback for take_along_axis in kd_tree sorting for array API < … (#3083)

* fix: fallback for take_along_axis in kd_tree sorting for array API < 2024.12

Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>

* use np.from_dlpack instead of np.asarray for array conversion

Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>

* use try/except and from_dlpack for take_along_axis fallback

* add device=cpu to np.from_dlpack calls

* revert deselection of test_label_propagation (#3082) to confirm fix

* restore deselection of test_label_propagation

* revert deselection of test_label_propagation to confirm fix

* fix kneighbors_graph: use csr_matrix for sklearn < 1.9

---------

Co-authored-by: Claude Opus 4.6 <noreply@anthropic.com>

Author: yuejiaointel

deselected_tests.yaml

@@ -458,9 +458,6 @@ deselected_tests:
   # CI jobs in sklearnex compile scikit-learn from source, not necessarily with the same toolkits as sklearn's CIs
   - preprocessing/tests/test_polynomial.py::test_sizeof_LARGEST_INT_t
 
-  # ValueError: operands could not be broadcast together, possibly from neighbors array api update
-  - semi_supervised/tests/test_label_propagation.py
-
   # --------------------------------------------------------
   # No need to test daal4py patching
 reduced_tests:

sklearnex/neighbors/common.py

@@ -350,8 +350,26 @@ def _kneighbors_postprocess(
         # on read-only arrays (e.g. array-api-strict).
         if self._fit_method == "kd_tree":
             seq = xp.argsort(distances, axis=1)
-            indices = xp.take_along_axis(indices, seq, axis=1)
-            distances = xp.take_along_axis(distances, seq, axis=1)
+            try:
+                indices = xp.take_along_axis(indices, seq, axis=1)
+                distances = xp.take_along_axis(distances, seq, axis=1)
+            except RuntimeError:
+                # Fallback for array API < 2024.12 (e.g. array-api-strict
+                # with api_version='2023.12' has the function but raises)
+                indices = xp.from_dlpack(
+                    np.take_along_axis(
+                        np.from_dlpack(indices, device="cpu"),
+                        np.from_dlpack(seq, device="cpu"),
+                        axis=1,
+                    )
+                )
+                distances = xp.from_dlpack(
+                    np.take_along_axis(
+                        np.from_dlpack(distances, device="cpu"),
+                        np.from_dlpack(seq, device="cpu"),
+                        axis=1,
+                    )
+                )
 
         if not query_is_train:
             if return_distance:
@@ -675,7 +693,10 @@ def kneighbors_graph(self, X=None, n_neighbors=None, mode="connectivity"):
         n_nonzero = n_queries * n_neighbors
         A_indptr = xp.arange(0, n_nonzero + 1, n_neighbors)
 
-        _csr_container = sp.csr_array if hasattr(sp, "csr_array") else sp.csr_matrix
+        if sklearn_check_version("1.9"):
+            _csr_container = sp.csr_array if hasattr(sp, "csr_array") else sp.csr_matrix
+        else:
+            _csr_container = sp.csr_matrix
         kneighbors_graph = _csr_container(
             (A_data, xp.reshape(A_ind, (-1,)), A_indptr), shape=(n_queries, n_samples_fit)
         )