Fix output channel calculation in PixelShuffle2DUpBlock

[wli51]

Older behavior is erroneous behavior from previous refactor due to lack of caution.

With pixel shuffle up sampling, the channel should reduce proportional to the combined 2D spatial expansion as opposed to remaining identical:

# Conceptually reasonable upsampling via pixel shuffling:
C × H × W
   ↓ Conv: C → C
C × H × W
   ↓ PixelShuffle(2)
C/4 × 2H × 2W # information is moved from channel to result in 4x increase in pixel along H&W, so channel loses 4x

# Earlier implementation
C × H × W
   ↓ Conv: C → 4C
4C × H × W
   ↓ PixelShuffle(2)
C × 2H × 2W # C remains unchanged despite pixels are shuffled from C into H&W, 
# because the pixel shuffle operation is itself rigid, this implementation effectively forces the conv2d layers bundled in
# `PixelShuffle` to invent 4x information. 
# Although this boosts model capacity because there are extra learnable parameters I prefer to have no  

[Copilot]

Pull request overview

Corrects PixelShuffle upsampling channel sizing while adding an opt-in legacy preservation mode.

Changes:

• Reduces output channels proportionally to 2D expansion.
• Adds channel-preservation configuration.
• Updates stage calculations and tests.

Reviewed changes

Copilot reviewed 4 out of 4 changed files in this pull request and generated 1 comment.

File	Summary
tests/models/test_up_down_blocks.py	Tests reduced and preserved channel behavior.
src/virtual_stain_flow/models/unext.py	Wires and serializes channel preservation; missing-key fallback may break loading older checkpoints.
src/virtual_stain_flow/models/stages.py	Updates stage spatial-dimension calculations.
src/virtual_stain_flow/models/blocks/up_down_blocks.py	Implements proportional PixelShuffle channel reduction.

---

💡 Add a code-review agent skill or configure MCP servers for context-aware, tailored reviews. Learn more in the docs.