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

Code before:

    # test callable kernel
    trans = Nystroem(n_components=2, kernel=_linear_kernel, random_state=rnd)
    X_transformed = trans.fit(X).transform(X)
    assert X_transformed.shape == (X.shape[0], 2)

    # test that available kernels fit and transform
    kernels_available = kernel_metrics()
    for kern in kernels_available:
        if (
            kern != "laplacian"
        ):  # Laplacian kernel not supported for non-numpy namespaces
            trans = Nystroem(n_components=2, kernel=kern, random_state=rnd)
            X_transformed = trans.fit(X).transform(X)
            assert X_transformed.shape == (X.shape[0], 2)


@pytest.mark.parametrize(
    "array_namespace, device, dtype_name", yield_namespace_device_dtype_combinations()
)
def test_nystroem_approximation_array_api(array_namespace, device, dtype_name):
    xp = _array_api_for_tests(array_namespace, device)
    # some basic tests
    rnd = np.random.RandomState(0)
    X_np = rnd.uniform(size=(10, 4))
    X_xp = xp.asarray(X_np)

    with config_context(array_api_dispatch=True):
        # With n_components = n_samples this is exact
        X_xp_transformed = Nystroem(n_components=X_xp.shape[0]).fit_transform(X_xp)
        X_xp_transformed_np = _convert_to_numpy(X_xp_transformed, xp=xp)
        K = rbf_kernel(X_np)
        assert_array_almost_equal(np.dot(X_xp_transformed_np, X_xp_transformed_np.T), K)

        trans = Nystroem(n_components=2, random_state=rnd)
        X_xp_transformed = trans.fit(X_xp).transform(X_xp)
        assert X_xp_transformed.shape == (X_xp.shape[0], 2)

        # test callable kernel
        trans = Nystroem(n_components=2, kernel=_linear_kernel, random_state=rnd)
        X_xp_transformed = trans.fit(X_xp).transform(X_xp)
        assert X_xp_transformed.shape == (X_xp.shape[0], 2)

        # test that available kernels fit and transform
        kernels_available = kernel_metrics()
        for kern in kernels_available:
            trans = Nystroem(n_components=2, kernel=kern, random_state=rnd)
            X_xp_transformed = trans.fit(X_xp).transform(X_xp)
            assert X_xp_transformed.shape == (X_xp.shape[0], 2)



Review comment (suggestion):
Please update all the tests to make sure that we actually use the `device` returned by `yield_namespace_device_dtype_combinations`:

```suggestion
    X_xp = xp.asarray(X_np, device=device)
```

Code after:

    # test callable kernel
    trans = Nystroem(n_components=2, kernel=_linear_kernel, random_state=rnd)
    X_transformed = trans.fit(X).transform(X)
    assert X_transformed.shape == (X.shape[0], 2)

    # test that available kernels fit and transform
    kernels_available = kernel_metrics()
    for kern in kernels_available:
        trans = Nystroem(n_components=2, kernel=kern, random_state=rnd)
        X_transformed = trans.fit(X).transform(X)
        assert X_transformed.shape == (X.shape[0], 2)


@pytest.mark.parametrize(
    "array_namespace, device, dtype_name", yield_namespace_device_dtype_combinations()
)
@pytest.mark.parametrize(
    "kernel", list(kernel_metrics()) + [_linear_kernel, "precomputed"]
)
@pytest.mark.parametrize("n_components", [2, 100])
def test_nystroem_approximation_array_api(
    array_namespace, device, dtype_name, kernel, n_components
):
    xp = _array_api_for_tests(array_namespace, device)
    rnd = np.random.RandomState(0)
    n_samples = 10
    # Ensure full-rank linear kernel to limit the impact of device-specific
    # rounding discrepancies.
    n_features = 2 * n_samples
    X_np = rnd.uniform(size=(n_samples, n_features)).astype(dtype_name)
    if kernel == "precomputed":
        X_np = rbf_kernel(X_np[:n_components])

    X_xp = xp.asarray(X_np, device=device)

    nystroem = Nystroem(n_components=n_components, kernel=kernel, random_state=0)
    X_np_transformed = nystroem.fit_transform(X_np)

    with config_context(array_api_dispatch=True):
        X_xp_transformed = nystroem.fit_transform(X_xp)
        X_xp_transformed_np = _convert_to_numpy(X_xp_transformed, xp=xp)

        for attribute_name in ["components_", "normalization_"]:
            xp_attr, _, device_attr = get_namespace_and_device(
                getattr(nystroem, attribute_name)
            )
            assert xp_attr is xp
            assert device_attr == array_device(X_xp)

    atol = _atol_for_type(dtype_name)