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

Code before:
        # For this backend, we can directly use "sqeuclidean".
        if metric_kwargs.get("squared", False) and metric == "euclidean":
            metric = "sqeuclidean"
            metric_kwargs = {}

        indices = ArgKmin.compute(
            X=X,
            Y=Y,
            k=1,
            metric=metric,
            metric_kwargs=metric_kwargs,
            strategy="auto",
            return_distance=False,
        )
        indices = indices.flatten()
    else:
        # Joblib-based backend, which is used when user-defined callable
        # are passed for metric.

        # This won't be used in the future once PairwiseDistancesReductions support:
        #   - DistanceMetrics which work on supposedly binary data
        #   - CSR-dense and dense-CSR case if 'euclidean' in metric.

        # Turn off check for finiteness because this is costly and because arrays
        # have already been validated.
        with config_context(assume_finite=True):
            indices_list = list(
                # This returns an np.ndarray generator whose arrays we need
                # to flatten into one.
                pairwise_distances_chunked(
                    X, Y, reduce_func=_argmin_reduce, metric=metric, **metric_kwargs
                )
            )
            indices = xp.concat(indices_list, axis=0)

    return indices


@validate_params(
    {"X": ["array-like", "sparse matrix"], "Y": ["array-like", "sparse matrix", None]},
    prefer_skip_nested_validation=True,
)
def haversine_distances(X, Y=None):
    """Compute the Haversine distance between samples in X and Y.

    The Haversine (or great circle) distance is the angular distance between
    two points on the surface of a sphere. The first coordinate of each point
    is assumed to be the latitude, the second is the longitude, given
    in radians. The dimension of the data must be 2.

    .. math::

Review comment (suggestion):
Instead of the diff (adding a new variable `indices_list`), you could just change `np.concatenate` to `xp.concat`. I think.

Code after:
        # This is an adaptor for one "sqeuclidean" specification.
        # For this backend, we can directly use "sqeuclidean".
        if metric_kwargs.get("squared", False) and metric == "euclidean":
            metric = "sqeuclidean"
            metric_kwargs = {}

        indices = ArgKmin.compute(
            X=X,
            Y=Y,
            k=1,
            metric=metric,
            metric_kwargs=metric_kwargs,
            strategy="auto",
            return_distance=False,
        )
        indices = indices.flatten()
    else:
        # Joblib-based backend, which is used when user-defined callable
        # are passed for metric.

        # This won't be used in the future once PairwiseDistancesReductions support:
        #   - DistanceMetrics which work on supposedly binary data
        #   - CSR-dense and dense-CSR case if 'euclidean' in metric.

        # Turn off check for finiteness because this is costly and because arrays
        # have already been validated.
        with config_context(assume_finite=True):
            indices = xp.concat(
                list(
                    pairwise_distances_chunked(
                        X, Y, reduce_func=_argmin_reduce, metric=metric, **metric_kwargs
                    )
                ),
                axis=0,
            )

    return indices


@validate_params(
    {"X": ["array-like", "sparse matrix"], "Y": ["array-like", "sparse matrix", None]},
    prefer_skip_nested_validation=True,
)
def haversine_distances(X, Y=None):
    """Compute the Haversine distance between samples in X and Y.

    The Haversine (or great circle) distance is the angular distance between
    two points on the surface of a sphere. The first coordinate of each point
    is assumed to be the latitude, the second is the longitude, given
    in radians. The dimension of the data must be 2.