fix(loss): make ChunkedCrossEntropy reductions independent of chunk_len

[kabirvashisht4-glitch]

What does this PR do ?

Makes ChunkedCrossEntropy return the same loss regardless of chunk_len.
chunk_len is a memory knob; it was changing the answer.

The fallback loop accumulated per-chunk results with +=
(components/loss/chunked_ce.py L199-216). That is valid for
reduction="sum" — which takes the separate _ChunkedCrossEntropySum kernel
anyway — but wrong for the two reductions that actually reach the loop:

• "mean" summed per-chunk means, returning num_chunks x the loss.
• "none" element-wise added the per-chunk vectors instead of
  concatenating, returning chunk_len values instead of one per token.

Neither raised. Measured against F.cross_entropy on 4096 tokens:

chunk_len	chunks	"mean" before	"mean" after
32 (default)	128	128.0x	1.0x
1024	4	4.0x	1.0x
2048	2	2.0x	1.0x
4096	1	1.0x	1.0x

"none" returned (32,) instead of (4096,) at the default chunk_len.

Note the last row: with chunk_len >= seq_len there is exactly one chunk, so
accumulation is a no-op and the result is correct. The error only appears once
the sequence is long enough to chunk — which is the case this loss exists for —
so a short smoke test passes.

Because the gradients scale with the loss, a config using reduction: mean
trained with an effectively multiplied learning rate.

Changelog

• "mean" accumulates the per-chunk sum and divides once by the
  non-ignored token count, which matches F.cross_entropy(..., reduction="mean")
  exactly rather than approximately.
• "none" concatenates the per-chunk vectors.
• "sum" is untouched — it never used this loop.
• New TestReductionIsChunkInvariant in tests/unit_tests/loss/test_chunked_ce.py
  comparing all three reductions against F.cross_entropy across
  chunk_len in {1, 32, 96, 100, 128, 512} — including 96 and 100, which do
  not divide the 512-token sequence evenly — with and without ignored positions
  scattered across chunk boundaries.

Before your PR is "Ready for review"

Pre checks:

• Make sure you read and followed Contributor guidelines
• Did you write any new necessary tests?
• Did you add or update any necessary documentation?

The class docstring already describes the fallback loop; its behaviour now
matches the description, so no doc change was needed.

Verification (CPU, no GPU needed):

• Reverting the source hunk fails 15 of the new cases (every "mean" and
  "none" case with more than one chunk); the "sum" cases and the
  single-chunk cases pass either way, which is intended — they pin the paths
  this must not change.
• pytest tests/unit_tests/loss/ → 318 passed, 36 skipped, no regressions.
• ruff format / ruff check clean.

Scope note. "mean" over a batch where every label is ignored divides by
zero and yields NaN, exactly as F.cross_entropy does — I kept that identical
rather than changing it here, since this PR is only about chunk-invariance. That
separate question is #3796 / #3797.

Additional Information

• Related to ChunkedCrossEntropy returns num_chunks x the loss for reduction='mean' and wrong-shaped output for 'none' #3803

[copy-pr-bot]

This pull request requires additional validation before any workflows can run on NVIDIA's runners.

Pull request vetters can view their responsibilities here.

Contributors can view more details about this message here.