def test_cluster_count_intra_class_distances(
    features: torch.Tensor, expected: torch.Tensor
) -> None:
    """
    Test _count_intra_class_distances method of HardClusterSampler.

    Args:
        features: tensor of shape (p, k, embed_dim), where p is a number of
        classes in the batch, k is a number of samples for each class,
        embed_dim is an embedding size -- features grouped by labels
        expected: tensor of shape (p, k) -- expected distances from mean
        vectors of classes to corresponding features
    """
    sampler = HardClusterSampler()
    mean_vectors = features.mean(1)
    distances = sampler._count_intra_class_distances(features, mean_vectors)
    assert (distances == expected).all()

    Test _count_intra_class_distances method of HardClusterSampler.

    Args:
        features: tensor of shape (p, k, embed_dim), where p is a number of
        classes in the batch, k is a number of samples for each class,
        embed_dim is an embedding size -- features grouped by labels
        expected: tensor of shape (p, k) -- expected distances from mean
        vectors of classes to corresponding features
    
tests/contrib/data/test_sampler_inbatch.py