🐛[BUG]: Variance in metrics.general.ensemble_metrics ignores the dim argument #2003

Version

main at ff5d19d

On which installation method(s) does this occur?

Source

Describe the issue

physicsnemo.metrics.general.ensemble_metrics.Variance.__call__ accepts a dim argument but two lines inside it still assume the ensemble dimension is the leading one.

1. The sample count is taken from inputs.shape[0] instead of inputs.shape[dim]. Mean.__call__ a few lines above uses inputs.shape[dim].
2. The centered sum of squares is torch.sum((inputs - self.sum / self.n) ** 2, dim=dim). self.sum has the reduced shape, so the subtraction only broadcasts when dim is the leading dimension.

As a result, for any dim other than 0 the call either raises a broadcasting RuntimeError, or, when the size of dim happens to equal the size of the leading dimension, it runs through and silently returns a wrong variance. The silent case is the worrying one because it is the shape you get from a square ensemble by time grid.

_update_var has the same broadcasting problem for a batch_dim other than 0 because it computes inputs - temp_sum / temp_n with the already reduced temp_sum.

The existing test_means_var only exercises dim=0 and is skipped when CUDA is not available, so CI never saw this.

I expected Variance(...)(x, dim=d) to match torch.var(x, dim=d) for every d, the same way Mean(...)(x, dim=d) matches torch.mean(x, dim=d).

Minimum reproducible example

import torch
import physicsnemo.metrics.general.ensemble_metrics as em

# Broadcasting error
x = torch.randn(4, 6, 5)
em.Variance((4, 5))(x, dim=1)
# RuntimeError: The size of tensor a (6) must match the size of tensor b (4) at non-singleton dimension 1

# Silent wrong result when the sizes happen to match
x = torch.randn(6, 6, 5)
var = em.Variance((6, 5))(x, dim=1)
print(torch.allclose(var, torch.var(x, dim=1)))
# False

# Mean is fine on the same input
print(torch.allclose(em.Mean((6, 5))(x, dim=1), torch.mean(x, dim=1)))
# True

Relevant log output

RuntimeError: The size of tensor a (6) must match the size of tensor b (4) at non-singleton dimension 1

Environment details

Bare-metal, CPU only, Python 3.12, torch CPU wheel, physicsnemo installed from source with pip install -e .

I have a fix ready with a CPU test that covers dim 1, 2 and -1 for both the error case and the silent case, and I will open a PR that references this issue.