docs: add specific tensor shape annotations to metrics.py

[Junaid1206]

Summary

This PR adds comprehensive Google-style docstrings and explicit tensor shape annotations to the metrics.py module.

Changes

• Documented all metric functions: wmse, wmae, nll, and crps_gauss.
• Added explicit tensor shape tracking (e.g., (..., N, d_state)) to clarify broadcasting behavior and batch dimensions.
• Expanded documentation of mask_and_reduce_metric to clearly explain spatial masking and grid/variable reduction logic.

Motivation and Context

Following the discussion in #252, this PR addresses a documentation gap in the evaluation module.

The evaluation logic is central to understanding how model performance is measured, particularly with respect to:

• Probabilistic metrics
• Weighted reductions
• Masked spatial evaluation

This PR focuses purely on documentation improvements and does not modify functionality.

Dependencies

None.

Type of Change

• Bug fix (non-breaking change that fixes an issue)
• New feature (non-breaking change that adds functionality)
• Breaking change (fix or feature that would cause existing functionality to not work as expected)
• Documentation (addition or improvement to documentation)

Checklist Before Requesting Review

• My branch is up-to-date with the target branch
• I have performed a self-review of my code
• For any new/modified functions/classes, I have added docstrings describing purpose, inputs, and return values
• I have added inline comments where necessary
• I have added tests (not applicable – documentation-only PR)
• I have given the PR a clear, imperative title
• I have requested a reviewer and assigned @sadamov

Author Checklist After Review

• I have added a line to the CHANGELOG (if required)

Checklist for Assignee

• PR is up to date with the base branch
• Tests pass
• (If applicable) Assigned to the appropriate milestone

[sadamov]

@Junaid1206 thank you, could you please revert to the original template. The current one is missing some bits.

[Junaid1206]

Updated the description to include the full original template with all checklist sections. Thanks for pointing that out, @sadamov

[joeloskarsson]

The description still looks very strange, could you have another look at the template.

[Junaid1206]

Hi @joeloskarsson, I apologize for the formatting issues in the previous update. I have now carefully re-applied the original repository template, ensuring all Markdown sections (headings, checklists, and bullet points) are correctly formatted and rendered.

Could you please take a look and let me know if this version meets the project's standards? Thank you for your patience.

Best regards
Junaid

[Junaid1206]

Hi @sadamov @joeloskarsson, I hope you're doing well. I have updated the PR description and fixed the formatting issues as per the repository template a few days ago. Please let me know if any further changes are required or if this is ready for review. Thank you!

Best Regards,
Junaid shah

[joeloskarsson]

To now consolidate all docstring improvements in one place, I think it would be best to add these as suggestions there (I think you already started looking at this :)). So I'll close this in favor of that workflow.