[BUG] Fix dimension assumption in QuantileLoss.loss for 2D inputs

[Kundan-CR7]

Reference Issues/PRs

Fixes #1795

What does this implement/fix? Explain your changes.

This PR fixes a dimension assumption in QuantileLoss.loss.

The current implementation concatenates per-quantile losses using:

torch.cat(losses, dim=2)

This assumes that y_pred is always 3D (e.g., [batch, time, quantiles]).
However, when y_pred is 2D (e.g., [batch, quantiles]), this raises:

IndexError: Dimension out of range

To make the implementation dimension-safe, concatenation has been changed to:

torch.cat(losses, dim=-1)

Since the quantile dimension is always the last dimension of y_pred, using dim=-1 ensures compatibility with both:

• 2D inputs: [batch, quantiles]
• 3D inputs: [batch, time, quantiles]

This resolves the issue without altering the loss computation logic.

What should a reviewer concentrate their feedback on?

• Correctness of replacing dim=2 with dim=-1
• Consistency with expected quantile dimension semantics
• Regression test coverage for 2D inputs
• Backward compatibility with existing 3D use cases

Did you add any tests for the change?

Yes.

Added regression tests in test_quantile_loss.py to verify:

• 2D input no longer raises IndexError
• 3D input behavior remains unchanged
• Output shapes are correct for both cases

All existing tests pass locally:

862 passed, 2 skipped

Any other comments?

This change only removes the fixed-dimension assumption and makes concatenation robust to varying input dimensionality.
No changes were made to the quantile loss computation itself.

PR checklist

• The PR title starts with either [ENH], [MNT], [DOC], or [BUG].
• Added/modified tests
• Used pre-commit hooks when committing to ensure that code is compliant with hooks.

[codecov]

Codecov Report

✅ All modified and coverable lines are covered by tests.
⚠️ Please upload report for BASE (main@602ee15). Learn more about missing BASE report.

Additional details and impacted files

@@           Coverage Diff           @@
##             main    #2113   +/-   ##
=======================================
  Coverage        ?   86.63%           
=======================================
  Files           ?      165           
  Lines           ?     9732           
  Branches        ?        0           
=======================================
  Hits            ?     8431           
  Misses          ?     1301           
  Partials        ?        0           

Flag	Coverage Δ
cpu	86.63% <100.00%> (?)
pytest	86.63% <100.00%> (?)

Flags with carried forward coverage won't be shown. Click here to find out more.

☔ View full report in Codecov by Sentry.
📢 Have feedback on the report? Share it here.

🚀 New features to boost your workflow:

• ❄️ Test Analytics: Detect flaky tests, report on failures, and find test suite problems.

[PranavBhatP]

@Kundan-CR7 thanks for the contribution! Your change is noted, and I think it would make sense to have dim=-1 instead of dim=2. But the reason why dim=2 was because it was ensured that we always have the input as [batch, time, quantiles] into the torch.cat function. That is my understanding, of course your change still stands.

[PranavBhatP]

Add separate units tests for metrics in the already existing ..\tests\test_metrics.py file.

[Kundan-CR7]

@PranavBhatP Thanks a lot for the clarification and the quick review!

That makes sense — I understand the original assumption about the input always being [batch, time, quantiles]. I’ve moved the regression test into the existing tests/test_metrics.py file as suggested to keep the test structure consistent.

This is my first PR in this repository, so I really appreciate the feedback and guidance.

[phoeenniixx]

I have a doubt:
should the QuantileLoss support 2D input? I mean we always assume it to be 3D no?