Repository: scikit-learn/scikit-learn
File: sklearn/metrics/cluster/_unsupervised.py (Python)

Code before:
        Predicted labels for each sample.

    Returns
    -------
    score : float
        The resulting Calinski-Harabasz score.

    References
    ----------
    .. [1] `T. Calinski and J. Harabasz, 1974. "A dendrite method for cluster
       analysis". Communications in Statistics
       <https://www.tandfonline.com/doi/abs/10.1080/03610927408827101>`_

    Examples
    --------
    >>> from sklearn.datasets import make_blobs
    >>> from sklearn.cluster import KMeans
    >>> from sklearn.metrics import calinski_harabasz_score
    >>> X, _ = make_blobs(random_state=0)
    >>> kmeans = KMeans(n_clusters=3, random_state=0,).fit(X)
    >>> calinski_harabasz_score(X, kmeans.labels_)
    114.8...
    """

    xp, _, device_ = get_namespace_and_device(X, labels)
    X = xp.asarray(X, dtype=_max_precision_float_dtype(xp, device_), device=device_)
    X, labels = check_X_y(X, labels)
    le = LabelEncoder()
    labels = le.fit_transform(labels)

    n_samples, _ = X.shape
    n_labels = le.classes_.shape[0]

    check_number_of_labels(n_labels, n_samples)

    extra_disp, intra_disp = 0.0, 0.0
    mean = xp.mean(X, axis=0)
    for k in range(n_labels):
        cluster_k = X[labels == k]
        mean_k = xp.mean(cluster_k, axis=0)
        extra_disp += cluster_k.shape[0] * xp.sum((mean_k - mean) ** 2)
        intra_disp += xp.sum((cluster_k - mean_k) ** 2)

    return float(
        1.0
        if intra_disp == 0.0
        else extra_disp * (n_samples - n_labels) / (intra_disp * (n_labels - 1.0))
    )


@validate_params(

Review comment (suggestion):
The issue with not being able to use `astype` is because for in the common tests for metrics we also test for a case where `X` and `labels` can be tensors but `array_api_dispatch=False`. This results in `xp` being equal to `numpy` and that creates issues when calling `numpy.astype` on a `Tensor`. The following code should fix this. It will hopefully not create a copy when we are dealing with the same `namespace`. Also when we are dealing with this discrepancy of `xp=numpy` and `X and labels being tensors` we might be able to simply convert to `numpy`.
```suggestion
    xp, _, device_ = get_namespace_and_device(X, labels)
    X = xp.asarray(X, device=device_)
    X = xp.astype(X, _max_precision_float_dtype(xp, device_), copy=False)
```


Code after:
        to a single data point.

    labels : array-like of shape (n_samples,)
        Predicted labels for each sample.

    Returns
    -------
    score : float
        The resulting Calinski-Harabasz score.

    References
    ----------
    .. [1] `T. Calinski and J. Harabasz, 1974. "A dendrite method for cluster
       analysis". Communications in Statistics
       <https://www.tandfonline.com/doi/abs/10.1080/03610927408827101>`_

    Examples
    --------
    >>> from sklearn.datasets import make_blobs
    >>> from sklearn.cluster import KMeans
    >>> from sklearn.metrics import calinski_harabasz_score
    >>> X, _ = make_blobs(random_state=0)
    >>> kmeans = KMeans(n_clusters=3, random_state=0,).fit(X)
    >>> calinski_harabasz_score(X, kmeans.labels_)
    114.8...
    """

    xp, _, device_ = get_namespace_and_device(X, labels)

    if _is_numpy_namespace(xp) and not is_numpy_array(X):
        # This is required to handle the case where `array_api_dispatch` is False but
        # we are still dealing with `X` as a non-NumPy array e.g. a PyTorch tensor.
        X = _convert_to_numpy(X, xp=xp)
    else:
        X = xp.astype(X, _max_precision_float_dtype(xp, device_), copy=False)
    X, labels = check_X_y(X, labels)
    le = LabelEncoder()
    labels = le.fit_transform(labels)

    n_samples, _ = X.shape
    n_labels = le.classes_.shape[0]

    check_number_of_labels(n_labels, n_samples)

    extra_disp, intra_disp = 0.0, 0.0
    mean = xp.mean(X, axis=0)
    for k in range(n_labels):
        cluster_k = X[labels == k]
        mean_k = xp.mean(cluster_k, axis=0)
        extra_disp += cluster_k.shape[0] * xp.sum((mean_k - mean) ** 2)
        intra_disp += xp.sum((cluster_k - mean_k) ** 2)