fix(callbacks): wrap visualisation forwards in the precision plugin's autocast context

[dwromero]

Summary

Visualization callbacks invoke the model outside Lightning's training/validation step precision context. In mixed-precision runs, the logged predictions can therefore differ from the predictions used to calculate validation loss.

Wrap the forward in trainer.precision_plugin.forward_context() in the sequence, image and volume visualization callbacks, then convert predictions to float32 for NumPy/matplotlib. Standard FP32 precision uses the plugin's no-op context.

Rebased onto main at dcf169ca02bb0aeba5efbf50363fbf4ea123abe1. The production change is limited to three forward sites across two existing files: 12 added lines and 6 removed lines. The refreshed commits are signed off, credit David W. Romero's original fix, and keep regression tests in a separate commit.

Validation

• python -m pytest tests/test_visualization_precision.py tests/test_image_grid_callback.py -q -o addopts='' — 22 passed.
• Six new CPU cases exercise complete sequence/image/volume rendering under FP32 and BF16, compare callback predictions exactly with a precision-plugin reference, and verify logging completes.
• All three BF16 regression cases fail on unmodified current main.
• All repository pre-commit hooks passed on the changed files.
• CUDA/FP16 execution and the historical checkpoint metrics from the original PR description were not revalidated; those historical measurements are not acceptance evidence for this revision.

Summary by CodeRabbit

• Bug Fixes

  • Validation visualizations now use the configured precision context during model inference.
  • Visualization outputs are consistently converted for reliable image, volume, and sequence rendering.
  • Prevented precision-related inconsistencies between validation results and logged visualizations.

• Tests

  • Added coverage for FP32 and BF16 mixed-precision visualization workflows.
  • Verified prediction accuracy, rendering completion, precision-state cleanup, and image logging.

[coderabbitai]

Important

Review skipped

Auto reviews are disabled on this repository. Please check the settings in the CodeRabbit UI or the .coderabbit.yaml file in this repository. To trigger a single review, invoke the @coderabbitai review command.

⚙️ Run configuration

Configuration used: Path: .coderabbit.yaml

Review profile: CHILL

Plan: Advanced

Run ID: 73c45b9d-fcaa-45e2-9d21-d9fd91c8657b

You can disable this status message by setting the reviews.review_status to false in the CodeRabbit configuration file.

Use the checkbox below for a quick retry:

• 🔍 Trigger review

No actionable comments were generated in the recent review. 🎉

ℹ️ Recent review info

⚙️ Run configuration

Configuration used: Path: .coderabbit.yaml

Review profile: CHILL

Plan: Enterprise

Run ID: 847a7939-93e0-4610-a268-6cbeb6d08f13

📥 Commits

Reviewing files that changed from the base of the PR and between dcf169c and aedecfa.

📒 Files selected for processing (3)

• experiments/callbacks/image_grid_val_visualization.py
• experiments/callbacks/sequence_visualization_1d.py
• tests/test_visualization_precision.py

Included review availability: Your plan provides up to 12 included reviews per hour; 11 remain after this review.

---

📝 Walkthrough

Walkthrough

Changes

Visualization precision handling

Layer / File(s)	Summary
Callback precision integration experiments/callbacks/image_grid_val_visualization.py, experiments/callbacks/sequence_visualization_1d.py	Image, volume, and sequence callbacks run inference inside trainer.precision_plugin.forward_context(). They convert predictions to float before NumPy or Matplotlib processing.
Precision regression tests tests/test_visualization_precision.py	Tests cover image, volume, and sequence visualization in FP32 and BF16 modes. They verify precision state, prediction dtype and values, rendering, and image logging.

Priority: ⬇️ Low

Estimated code review effort: 3 (Moderate) | ~20 minutes

Change: Bug fix

Merge Risk: ⚪ Minimal · up to aedec

No merge-blocking visualization precision risk was identified.

🚥 Pre-merge checks | ✅ 4 | ❌ 1

❌ Failed checks (1 warning)

Check name	Status	Explanation	Resolution
Docstring Coverage	⚠️ Warning	Docstring coverage is 50.00% which is insufficient. The required threshold is 80.00%. Docstring coverage is scoped to functions touched by this diff. Analyzed 6 functions across 3 files.	Write docstrings for the functions missing them to satisfy the coverage threshold.

✅ Passed checks (4 passed)

Check name	Status	Explanation
Title check	✅ Passed	The title clearly and concisely identifies the main change: wrapping visualization forwards in the precision plugin's autocast context.
Description check	✅ Passed	The description explains the problem, implementation, affected callbacks, tests, validation results, and known limitations. It does not reproduce the repository template sections for environment setup…
Linked Issues check	✅ Passed	Check skipped because no linked issues were found for this pull request.
Out of Scope Changes check	✅ Passed	Check skipped because no linked issues were found for this pull request.

✨ Finishing Touches

📝 Generate docstrings

• 🤖 Coding Agent task started for docstring generation.
• Create a new PR

🧪 Generate unit tests (beta)

• Commit to this branch
• Create a new PR

---

_{Comment @coderabbitai help to get the list of available commands.}

[coderabbitai]

🤖 Completed: Generate docstrings for PR #114 — View commit a73ea31

[farhadrgh]

LGTM