spearman_corrcoef and kendall_rank_corrcoef silently downcast float64 to float32

🐛 Bug

spearman_corrcoef and kendall_rank_corrcoef silently return float32 for float64 inputs. Every other regression metric preserves the input dtype, and kendall_rank_corrcoef itself preserves it for variant="c" but not for "a" or "b", so this looks unintended rather than deliberate.

To Reproduce

import torch
from torchmetrics.functional import pearson_corrcoef, spearman_corrcoef, kendall_rank_corrcoef

preds = torch.rand(100, dtype=torch.float64)
target = torch.rand(100, dtype=torch.float64)

print(pearson_corrcoef(preds, target).dtype)       # torch.float64
print(spearman_corrcoef(preds, target).dtype)      # torch.float32  <-- 
print(kendall_rank_corrcoef(preds, target).dtype)  # torch.float32  <-- 

The class-based SpearmanCorrCoef and KendallRankCorrCoef return float32 as well.

kendall_rank_corrcoef is inconsistent across its own variants:

variant	output dtype for float64 input
"a"	float32
"b" (default)	float32
"c"	float64

Surveying the functional regression metrics with float64 inputs, these two are the only ones that downcast:

metric	float64 in -> out
mean_absolute_error, mean_squared_error, mean_absolute_percentage_error, symmetric_mean_absolute_percentage_error, mean_squared_log_error, weighted_mean_absolute_percentage_error, r2_score, explained_variance, pearson_corrcoef, concordance_corrcoef, log_cosh_error, relative_squared_error, tweedie_deviance_score, minkowski_distance, normalized_root_mean_squared_error	float64
spearman_corrcoef	float32
kendall_rank_corrcoef	float32

Why it matters beyond the dtype label

Spearman is rank-based, so two identical strictly-increasing vectors must give exactly 1.0 regardless of the spacing between values. When the differences are below float32 resolution the result drifts:

import torch
from torchmetrics.functional import spearman_corrcoef, pearson_corrcoef

a = torch.tensor([1.0, 1.0 + 1e-10, 1.0 + 2e-10, 1.0 + 3e-10], dtype=torch.float64)

print(spearman_corrcoef(a, a).item())  # 0.9999992251396179
print(pearson_corrcoef(a, a).item())   # 0.9999999999999999

Users who deliberately opt into float64 — common when correlating values with a small dynamic range — silently lose that precision.

Cause

Both come from an integer intermediate being divided, which falls back to the global default dtype (float32) rather than the input dtype.

spearman_corrcoef — _rank_data in functional/regression/spearman.py:

sum_ranks = torch.zeros_like(uniq, dtype=torch.int32)
sum_ranks.scatter_add_(0, inv, rank.to(torch.int32))
mean_ranks = sum_ranks / counts   # int32 / int64 -> float32, regardless of data.dtype
return mean_ranks[inv]

The ranks therefore come back float32 and everything downstream inherits it.

kendall_rank_corrcoef — _get_metric_metadata in functional/regression/kendall.py:

n_total = torch.tensor(preds.shape[0], device=preds.device)   # int64, no dtype from preds

The concordant/discordant/tie counts are int64 too, so in _calculate_tau the divisions for variants "a" and "b" produce float32. Variant "c" escapes it only because it builds preds_unique with dtype=preds.dtype, which promotes the expression.

Expected behavior

The output dtype should follow the input, as the other regression metrics do — float64 in, float64 out. For float16/bfloat16 inputs, promoting the rank/count arithmetic to float32 is still appropriate, since ranks are integers that need more mantissa than float16 offers for larger inputs.

Environment

• TorchMetrics: master (commit 37a80809)
• PyTorch: 2.13.0+cpu, Python 3.13, Windows

I have a fix and regression tests ready and am happy to open a PR if this is welcome.

Start in functional/regression/spearman.py at _rank_data and functional/regression/kendall.py at _get_metric_metadata and _calculate_tau, then inspect the existing regression tests. Done means float64 inputs preserve float64 outputs for both metrics and Kendall variants, with regression coverage for the reported precision case and class-based metrics.