Fix Variance ensemble metric ignoring the dim argument

[Nicholas022400701]

PhysicsNeMo Pull Request

Description

closes #2003

Variance.__call__ in physicsnemo/metrics/general/ensemble_metrics.py accepts a dim argument but still assumed the ensemble dimension is the leading one in two places. It took the sample count from inputs.shape[0] while Mean.__call__ right above it uses inputs.shape[dim], and it computed the centered sum of squares as torch.sum((inputs - self.sum / self.n) ** 2, dim=dim) with self.sum already reduced, which only broadcasts when dim is 0. For any other dim the call either raised a broadcasting RuntimeError, or, when the size of dim happened to match the size of the leading dimension, it ran through and returned a wrong variance without any warning. _update_var had the same broadcasting problem for a batch_dim other than 0.

Changes:

• Variance.__call__ takes the sample count from inputs.shape[dim] and unsqueezes the reduced sum back along dim before subtracting it, in both the distributed branch and the plain branch.
• _update_var unsqueezes temp_sum along batch_dim the same way.
• New CPU test test/metrics/test_ensemble_metrics_dim.py covers dim 1, 2 and -1 for a shape that used to raise and for a square shape that used to return a wrong value, and checks _update_var with the same batch_dim against torch.var on the concatenated input. The existing test_means_var only covers dim=0 and is skipped without CUDA, so this bug never showed up in CI.

Nothing else in the file changed. Variance.update still only supports the leading dimension, as it did before, because it has no dim argument.

Verification on a CPU GitHub Actions runner with torch CPU wheels and Python 3.12:

• Fix branch, 6 new cases pass plus ruff format --check, ruff check and interrogate on the touched files: https://github.com/Nicholas022400701/oss-ci/actions/runs/35351315842
• Same test file dropped onto upstream main, all 6 cases fail, 5 with the broadcasting RuntimeError and the square shape with a wrong value: https://github.com/Nicholas022400701/oss-ci/actions/runs/35350614463

I did not touch CHANGELOG.md in this PR. I can add a line under Fixed in the 2.3.0 section if you want it here.

AI disclosure: I used an AI coding agent to help write this patch, the tests and this description. I have read the change and the tests myself and I will answer review comments personally.

Checklist

• I am familiar with the Contributing Guidelines.
• New or existing tests cover these changes.
• The documentation is up to date with these changes. No public API or docstring changes, the dim argument now does what the docstring already says.
• The CHANGELOG.md is up to date with these changes.
• An issue is linked to this pull request.
• If I am implementing a new model or modifying any existing model, I have followed the Models Implementation Coding Standards. Not applicable, no model is touched.

Dependencies

None.

Review Process

All PRs are reviewed by the PhysicsNeMo team before merging.

Depending on which files are changed, GitHub may automatically assign a maintainer for review.

We are also testing AI-based code review tools (e.g., Greptile), which may add automated comments with a confidence score.
This score reflects the AI’s assessment of merge readiness and is not a qualitative judgment of your work, nor is
it an indication that the PR will be accepted / rejected.

AI-generated feedback should be reviewed critically for usefulness.
You are not required to respond to every AI comment, but they are intended to help both authors and reviewers.
Please react to Greptile comments with 👍 or 👎 to provide feedback on their accuracy.

[copy-pr-bot]

This pull request requires additional validation before any workflows can run on NVIDIA's runners.

Pull request vetters can view their responsibilities here.

Contributors can view more details about this message here.

[Copilot]

Copilot was unable to review this pull request because the user who requested the review has reached their quota limit.

[github-actions]

CODEOWNERS review map

Current for commit c8e89d4531c2. An approval covers every file listed for that owner; one owner is sufficient for shared files.

⏳ @dallasfoster — 1 file(s)

• physicsnemo/metrics/general/ensemble_metrics.py

⏳ @NickGeneva — 1 file(s)

• physicsnemo/metrics/general/ensemble_metrics.py

No CODEOWNER

• test/metrics/test_ensemble_metrics_dim.py

Comment /codeowners-info to refresh.

[greptile-apps]

The PR appears safe to merge; the dimension-aware variance calculations and regression tests are internally consistent.

Summary

This PR fixes variance calculations that previously ignored or incorrectly broadcast non-leading ensemble dimensions.

• Derives the sample count from the selected dimension.
• Restores the reduced axis before centering values in Variance.__call__ and _update_var.
• Adds CPU tests for positive and negative non-leading dimensions, including both prior broadcast-error and silent-wrong-result shapes.

_{Reviews (1) · Last reviewed commit: "Apply ruff format to the new test"}