Sum all video event dimensions in posterior KL divergence

[sylvesterkaczmarek]

Sum every non-batch dimension in DiagonalGaussianDistribution.kl(), including the channel axis of video posteriors [B, T, H, W, C]. Return one deterministic zero per batch item as well.

The hardcoded (1, 2, 3) reduction leaves a [B, C] result for video latents. Reducing that result with .mean() averages channels rather than summing the independent Gaussian contributions. In the 16-channel regression, the gradient is 0.03125 instead of 0.5, underweighting this KL term by 16. Both standard-normal KL and KL against another posterior are affected.

The four-dimensional image reduction is preserved. Posterior sampling, its mode, and inference decoding are unchanged. This intentionally changes video KL output from [B, C] to [B].

Validation

CPU, macOS arm64, Python 3.12, JAX 0.7.2, Flax 0.11.2:

• On unchanged 9912d29, four tests fail and the image-posterior control passes.
• All five tests pass after the fix: analytic per-example values, nonstandard reference Gaussian, JIT-compiled gradients, deterministic batch shape, and image compatibility.
• Built the wheel, installed it into the test environment, and reran all five tests outside the source checkout: passed.
• Ruff E9/F, Pyink on the new test, compilation, and git diff --check passed.

No pretrained-model training run or GPU execution was performed. The unpinned JAX 0.11.2/Flax 0.12.9 environment fails in Flax NNX import (HiPrimitive) even before the patch; this PR does not change dependencies.