fix(diffusion): broadcast transition sigma to replay mean rank #447

Problem

_transition_sigma() always returns a five-dimensional tensor with shape
[1, S, 1, 1, 1], regardless of the rank of the replayed transition means:

return sigma_t.reshape(1, -1, 1, 1, 1)

That shape only broadcasts correctly for image means shaped
[B, S, C, H, W]. The same helper is shared by FlowDPPO, FlowGRPO's reference
KL, and DiffusionOPD, while supported stages also return:

• video means: [B, S, C, T, H, W] (6-D)
• SenseNova packed-pixel means: [B, S, seq, C] (4-D)

Impact

For video means, KL computation fails with a broadcast error. For example:

means:   [2, 3, 16, 10, 64, 64]
sigma_t: [1, 3, 1, 1, 1]
RuntimeError: The size of tensor a (16) must match the size of tensor b (3)
              at non-singleton dimension 2

For four-dimensional SenseNova means, the common micro_batch_size: 1 case
silently broadcasts to a step cross-product instead of failing:

means:   [1, 3, 128, 3]
sigma_t: [1, 3, 1, 1, 1]
result:  [1, 3, 3, 128, 3]  # expected [1, 3, 128, 3]

With B > 1 and B != S, the same path raises a broadcast error instead.

This is reachable in the shipped
examples/diffusion/sensenova_u1_5/sensenova_u1_5_ocr_trainside_lora_kl_4x8.yaml:

• FlowGRPO
• beta: 0.01, which enables _reference_kl_loss
• micro_batch_size: 1
• three SDE steps

The silent extra S dimension changes both the reference-KL value and its
memory use.

Expected behavior

The transition standard deviation should broadcast only over the batch and
latent dimensions while preserving the step axis, for example:

sigma_t.reshape(1, S, *([1] * (means.ndim - 2)))

The exact API can either pass the mean rank into _transition_sigma() or shape
the returned per-step values at each mean consumer.

Suggested validation

Add focused checks for:

• image means [B, S, C, H, W]
• video means [B, S, C, T, H, W]
• packed means [B, S, seq, C]
• B=1 and B>1
• FlowDPPO, FlowGRPO reference KL, and DiffusionOPD callers

Context

Surfaced while reviewing #444; the bug predates that PR and is independent of
the replay-mean dtype fix.