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):
We need floating X for `xp.mean`, but I'm not sure if there is a cheaper way to do this.

When X is `torch.Tensor` we can't use `xp.astype`, though I imagine this is preferred:

```python
X = xp.astype(X, dtype=_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)