Skip to content

Commit dbbc3e0

Browse files
add torch examples
1 parent f21f741 commit dbbc3e0

File tree

3 files changed

+149
-71
lines changed

3 files changed

+149
-71
lines changed

doc/sources/array_api.rst

Lines changed: 96 additions & 46 deletions
Original file line numberDiff line numberDiff line change
@@ -110,52 +110,102 @@ Example usage
110110
GPU operations on GPU arrays
111111
----------------------------
112112

113-
.. code-block:: python
114-
115-
# Array API support from sklearn requires enabling it on SciPy too
116-
import os
117-
os.environ["SCIPY_ARRAY_API"] = "1"
118-
119-
import numpy as np
120-
import dpnp
121-
from sklearnex import config_context
122-
from sklearnex.linear_model import LinearRegression
123-
124-
# Random data for a regression problem
125-
rng = np.random.default_rng(seed=123)
126-
X_np = rng.standard_normal(size=(100, 10), dtype=np.float32)
127-
y_np = rng.standard_normal(size=100, dtype=np.float32)
128-
129-
# DPNP offers an array-API-compliant class where data can be on GPU
130-
X = dpnp.array(X_np, device="gpu")
131-
y = dpnp.array(y_np, device="gpu")
132-
133-
# Important to note again that array API must be enabled on scikit-learn
134-
model = LinearRegression()
135-
with config_context(array_api_dispatch=True):
136-
model.fit(X, y)
137-
138-
# Fitted attributes are now of the same class as inputs
139-
assert isinstance(model.coef_, X.__class__)
140-
141-
# Predictions are also of the same class
142-
with config_context(array_api_dispatch=True):
143-
pred = model.predict(X[:5])
144-
assert isinstance(pred, X.__class__)
145-
146-
# Fitted models can be passed array API inputs of a different class
147-
# than the training data, as long as their data resides in the same
148-
# device. This now fits a model using a non-NumPy class whose data is on CPU.
149-
X_cpu = dpnp.array(X_np, device="cpu")
150-
y_cpu = dpnp.array(y_np, device="cpu")
151-
model_cpu = LinearRegression()
152-
with config_context(array_api_dispatch=True):
153-
model_cpu.fit(X_cpu, y_cpu)
154-
pred_dpnp = model_cpu.predict(X_cpu[:5])
155-
pred_np = model_cpu.predict(X_cpu[:5].asnumpy())
156-
assert isinstance(pred_dpnp, X_cpu.__class__)
157-
assert isinstance(pred_np, np.ndarray)
158-
assert pred_dpnp.__class__ != pred_np.__class__
113+
.. tabs::
114+
.. tab:: With Torch tensors
115+
.. code-block:: python
116+
117+
# Array API support from sklearn requires enabling it on SciPy too
118+
import os
119+
os.environ["SCIPY_ARRAY_API"] = "1"
120+
121+
import numpy as np
122+
import torch
123+
from sklearnex import config_context
124+
from sklearnex.linear_model import LinearRegression
125+
126+
# Random data for a regression problem
127+
rng = np.random.default_rng(seed=123)
128+
X_np = rng.standard_normal(size=(100, 10), dtype=np.float32)
129+
y_np = rng.standard_normal(size=100, dtype=np.float32)
130+
131+
# Torch offers an array-API-compliant class where data can be on GPU (referred to as 'xpu')
132+
X = torch.tensor(X_np, device="xpu")
133+
y = torch.tensor(y_np, device="xpu")
134+
135+
# Important to note again that array API must be enabled on scikit-learn
136+
model = LinearRegression()
137+
with config_context(array_api_dispatch=True):
138+
model.fit(X, y)
139+
140+
# Fitted attributes are now of the same class as inputs
141+
assert isinstance(model.coef_, torch.Tensor)
142+
143+
# Predictions are also of the same class
144+
with config_context(array_api_dispatch=True):
145+
pred = model.predict(X[:5])
146+
assert isinstance(pred, torch.Tensor)
147+
148+
# Fitted models can be passed array API inputs of a different class
149+
# than the training data, as long as their data resides in the same
150+
# device. This now fits a model using a non-NumPy class whose data is on CPU.
151+
X_cpu = torch.tensor(X_np, device="cpu")
152+
y_cpu = torch.tensor(y_np, device="cpu")
153+
model_cpu = LinearRegression()
154+
with config_context(array_api_dispatch=True):
155+
model_cpu.fit(X_cpu, y_cpu)
156+
pred_torch = model_cpu.predict(X_cpu[:5])
157+
pred_np = model_cpu.predict(X_cpu[:5].numpy())
158+
assert isinstance(pred_torch, X_cpu.__class__)
159+
assert isinstance(pred_np, np.ndarray)
160+
assert pred_torch.__class__ != pred_np.__class__
161+
162+
.. tab:: With DPNP arrays
163+
.. code-block:: python
164+
165+
# Array API support from sklearn requires enabling it on SciPy too
166+
import os
167+
os.environ["SCIPY_ARRAY_API"] = "1"
168+
169+
import numpy as np
170+
import dpnp
171+
from sklearnex import config_context
172+
from sklearnex.linear_model import LinearRegression
173+
174+
# Random data for a regression problem
175+
rng = np.random.default_rng(seed=123)
176+
X_np = rng.standard_normal(size=(100, 10), dtype=np.float32)
177+
y_np = rng.standard_normal(size=100, dtype=np.float32)
178+
179+
# DPNP offers an array-API-compliant class where data can be on GPU
180+
X = dpnp.array(X_np, device="gpu")
181+
y = dpnp.array(y_np, device="gpu")
182+
183+
# Important to note again that array API must be enabled on scikit-learn
184+
model = LinearRegression()
185+
with config_context(array_api_dispatch=True):
186+
model.fit(X, y)
187+
188+
# Fitted attributes are now of the same class as inputs
189+
assert isinstance(model.coef_, X.__class__)
190+
191+
# Predictions are also of the same class
192+
with config_context(array_api_dispatch=True):
193+
pred = model.predict(X[:5])
194+
assert isinstance(pred, X.__class__)
195+
196+
# Fitted models can be passed array API inputs of a different class
197+
# than the training data, as long as their data resides in the same
198+
# device. This now fits a model using a non-NumPy class whose data is on CPU.
199+
X_cpu = dpnp.array(X_np, device="cpu")
200+
y_cpu = dpnp.array(y_np, device="cpu")
201+
model_cpu = LinearRegression()
202+
with config_context(array_api_dispatch=True):
203+
model_cpu.fit(X_cpu, y_cpu)
204+
pred_dpnp = model_cpu.predict(X_cpu[:5])
205+
pred_np = model_cpu.predict(X_cpu[:5].asnumpy())
206+
assert isinstance(pred_dpnp, X_cpu.__class__)
207+
assert isinstance(pred_np, np.ndarray)
208+
assert pred_dpnp.__class__ != pred_np.__class__
159209
160210
161211
``array-api-strict``

doc/sources/index.rst

Lines changed: 6 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -105,7 +105,7 @@ Note: executing on GPU has `additional system software requirements <https://www
105105
import os
106106
os.environ["SCIPY_ARRAY_API"] = "1"
107107
import numpy as np
108-
import dpnp
108+
import torch
109109
from sklearnex import patch_sklearn
110110
patch_sklearn()
111111
from sklearn import config_context
@@ -114,8 +114,8 @@ Note: executing on GPU has `additional system software requirements <https://www
114114
115115
X = np.array([[1., 2.], [2., 2.], [2., 3.],
116116
[8., 7.], [8., 8.], [25., 80.]], dtype=np.float32)
117-
X = dpnp.array(X, device="gpu")
118-
with config_context(array_api_dispatch=True)
117+
X = torch.tensor(X, device="xpu")
118+
with config_context(array_api_dispatch=True):
119119
clustering = DBSCAN(eps=3, min_samples=2).fit(X)
120120
121121
.. tab:: Without patching
@@ -138,14 +138,14 @@ Note: executing on GPU has `additional system software requirements <https://www
138138
import os
139139
os.environ["SCIPY_ARRAY_API"] = "1"
140140
import numpy as np
141-
import dpnp
141+
import torch
142142
from sklearnex import config_context
143143
from sklearnex.cluster import DBSCAN
144144
145145
X = np.array([[1., 2.], [2., 2.], [2., 3.],
146146
[8., 7.], [8., 8.], [25., 80.]], dtype=np.float32)
147-
X = dpnp.array(X, device="gpu")
148-
with config_context(array_api_dispatch=True)
147+
X = torch.tensor(X, device="xpu")
148+
with config_context(array_api_dispatch=True):
149149
clustering = DBSCAN(eps=3, min_samples=2).fit(X)
150150
151151

doc/sources/oneapi-gpu.rst

Lines changed: 47 additions & 19 deletions
Original file line numberDiff line numberDiff line change
@@ -139,30 +139,58 @@ This is particularly useful when multiple operations are performed on the same d
139139

140140
See :ref:`array_api` for details, instructions, and limitations. Example:
141141

142-
.. code-block:: python
142+
.. tabs::
143+
.. tab:: With Torch tensors
144+
.. code-block:: python
143145
144-
# Array API support from sklearn requires enabling it on SciPy too
145-
import os
146-
os.environ["SCIPY_ARRAY_API"] = "1"
146+
# Array API support from sklearn requires enabling it on SciPy too
147+
import os
148+
os.environ["SCIPY_ARRAY_API"] = "1"
147149
148-
import numpy as np
149-
import dpnp
150-
from sklearnex import config_context
151-
from sklearnex.linear_model import LinearRegression
150+
import numpy as np
151+
import torch
152+
from sklearnex import config_context
153+
from sklearnex.linear_model import LinearRegression
152154
153-
# Random data for a regression problem
154-
rng = np.random.default_rng(seed=123)
155-
X_np = rng.standard_normal(size=(100, 10), dtype=np.float32)
156-
y_np = rng.standard_normal(size=100, dtype=np.float32)
155+
# Random data for a regression problem
156+
rng = np.random.default_rng(seed=123)
157+
X_np = rng.standard_normal(size=(100, 10), dtype=np.float32)
158+
y_np = rng.standard_normal(size=100, dtype=np.float32)
157159
158-
# DPNP offers an array-API-compliant class where data can be on GPU
159-
X = dpnp.array(X_np, device="gpu")
160-
y = dpnp.array(y_np, device="gpu")
160+
# Torch offers an array-API-compliant class where data can be on GPU (referred to as 'xpu')
161+
X = torch.tensor(X_np, device="xpu")
162+
y = torch.tensor(y_np, device="xpu")
161163
162-
# Important to note again that array API must be enabled on scikit-learn
163-
model = LinearRegression()
164-
with config_context(array_api_dispatch=True):
165-
model.fit(X, y)
164+
# Important to note again that array API must be enabled on scikit-learn
165+
model = LinearRegression()
166+
with config_context(array_api_dispatch=True):
167+
model.fit(X, y)
168+
169+
.. tab:: With DPNP arrays
170+
.. code-block:: python
171+
172+
# Array API support from sklearn requires enabling it on SciPy too
173+
import os
174+
os.environ["SCIPY_ARRAY_API"] = "1"
175+
176+
import numpy as np
177+
import dpnp
178+
from sklearnex import config_context
179+
from sklearnex.linear_model import LinearRegression
180+
181+
# Random data for a regression problem
182+
rng = np.random.default_rng(seed=123)
183+
X_np = rng.standard_normal(size=(100, 10), dtype=np.float32)
184+
y_np = rng.standard_normal(size=100, dtype=np.float32)
185+
186+
# DPNP offers an array-API-compliant class where data can be on GPU
187+
X = dpnp.array(X_np, device="gpu")
188+
y = dpnp.array(y_np, device="gpu")
189+
190+
# Important to note again that array API must be enabled on scikit-learn
191+
model = LinearRegression()
192+
with config_context(array_api_dispatch=True):
193+
model.fit(X, y)
166194
167195
.. note::
168196
Not all estimator classes in the |sklearnex| support array API objects - see the list of :ref:`estimators with array API support <array_api_estimators>` for details.

0 commit comments

Comments
 (0)