Repository: scikit-learn/scikit-learn
File: sklearn/kernel_approximation.py (Python)

Code before:
        Returns
        -------
        self : object
            Returns the instance itself.
        """

        xp, _ = get_namespace(X)
        X = validate_data(self, X, accept_sparse="csr")
        rnd = check_random_state(self.random_state)
        n_samples = X.shape[0]

        # get basis vectors
        if self.n_components > n_samples:
            # XXX should we just bail?
            n_components = n_samples
            warnings.warn(
                "n_components > n_samples. This is not possible.\n"
                "n_components was set to n_samples, which results"
                " in inefficient evaluation of the full kernel."
            )

        else:
            n_components = self.n_components
        n_components = min(n_samples, n_components)
        inds = rnd.permutation(n_samples)
        basis_inds = xp.asarray(inds[:n_components], dtype=xp.int64)
        if sp.issparse(X):
            basis = X[basis_inds]
        else:
            basis = xp.take(X, basis_inds, axis=0)

        basis_kernel = pairwise_kernels(
            basis,
            metric=self.kernel,
            filter_params=True,
            n_jobs=self.n_jobs,
            **self._get_kernel_params(),
        )

        # sqrt of kernel matrix on basis vectors
        _, _, dtype = _find_floating_dtype_allow_sparse(basis_kernel, Y=None, xp=xp)
        basis_kernel = xp.asarray(basis_kernel, dtype=dtype)
        U, S, V = xp.linalg.svd(basis_kernel)
        S = xp.asarray(S, dtype=dtype)
        S = xp.maximum(S, xp.asarray(1e-12, dtype=dtype))
        U = xp.asarray(U, dtype=dtype)
        V = xp.asarray(V, dtype=dtype)
        self.normalization_ = xp.matmul(U / xp.sqrt(S), V)
        self.components_ = basis
        self.component_indices_ = basis_inds
        self._n_features_out = n_components

Review comment (suggestion):
and then always pass the device argument when calling `xp.asarray` to continue working on the same device.

Code after:
        basis_kernel = pairwise_kernels(
            basis,
            metric=self.kernel,
            filter_params=True,
            n_jobs=self.n_jobs,
            **self._get_kernel_params(),
        )

        # sqrt of kernel matrix on basis vectors
        _, _, dtype = _find_floating_dtype_allow_sparse(basis_kernel, Y=None, xp=xp)
        basis_kernel = xp.asarray(basis_kernel, dtype=dtype, device=device)
        U, S, V = xp.linalg.svd(basis_kernel)
        S = xp.clip(S, 1e-12, None)
        self.normalization_ = U / xp.sqrt(S) @ V
        self.components_ = basis
        self.component_indices_ = basis_inds
        self._n_features_out = n_components
        return self

    def transform(self, X):
        """Apply feature map to X.

        Computes an approximate feature map using the kernel
        between some training points and X.

        Parameters
        ----------
        X : array-like of shape (n_samples, n_features)
            Data to transform.

        Returns
        -------
        X_transformed : ndarray of shape (n_samples, n_components)
            Transformed data.
        """
        check_is_fitted(self)

        xp, _, device = get_namespace_and_device(X)
        X = validate_data(self, X, accept_sparse="csr", reset=False)

        kernel_params = self._get_kernel_params()
        embedded = pairwise_kernels(
            X,
            self.components_,
            metric=self.kernel,
            filter_params=True,
            n_jobs=self.n_jobs,
            **kernel_params,
        )
        dtype = _find_matching_floating_dtype(embedded, xp=xp)
        embedded = xp.asarray(embedded, dtype=dtype, device=device)