fix: forest dpctl predict queue misalignment (#2507)

ethanglaser · web-flow · commit d56958e5553f · 2025-06-26T00:13:42.000-07:00
* fix: forest dpctl predict queue misalignment

* black formatting

* sycl_queue -&gt; device
diff --git a/onedal/ensemble/forest.py b/onedal/ensemble/forest.py
@@ -522,7 +522,8 @@ def predict(self, X, queue=None):
 
         try:
             return xp.take(
-                xp.asarray(self.classes_), xp.astype(xp.reshape(pred, (-1,)), xp.int64)
+                xp.asarray(self.classes_, device=pred.sycl_queue),
+                xp.astype(xp.reshape(pred, (-1,)), xp.int64),
             )
         except AttributeError:
             return np.take(self.classes_, pred.ravel().astype(np.int64, casting="unsafe"))
diff --git a/sklearnex/ensemble/_forest.py b/sklearnex/ensemble/_forest.py
@@ -831,7 +831,8 @@ def _onedal_predict(self, X, queue=None):
         res = self._onedal_estimator.predict(X, queue=queue)
         try:
             return xp.take(
-                xp.asarray(self.classes_), xp.astype(xp.reshape(res, (-1,)), xp.int64)
+                xp.asarray(self.classes_, device=res.sycl_queue),
+                xp.astype(xp.reshape(res, (-1,)), xp.int64),
             )
         except AttributeError:
             return np.take(self.classes_, res.ravel().astype(np.int64, casting="unsafe"))

Original file line number	Diff line number	Diff line change
`@@ -522,7 +522,8 @@ def predict(self, X, queue=None):`
`522`	`522`
`523`	`523`	`try:`
`524`	`524`	`return xp.take(`
`525`		`- xp.asarray(self.classes_), xp.astype(xp.reshape(pred, (-1,)), xp.int64)`
	`525`	`+ xp.asarray(self.classes_, device=pred.sycl_queue),`
	`526`	`+ xp.astype(xp.reshape(pred, (-1,)), xp.int64),`
`526`	`527`	`)`
`527`	`528`	`except AttributeError:`
`528`	`529`	`return np.take(self.classes_, pred.ravel().astype(np.int64, casting="unsafe"))`