@@ -110,52 +110,102 @@ Example usage
110110GPU 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 ``
0 commit comments