fix(autograd): CrossEntropyLoss tape detach, softmax/variance backward on rank>=3

[michalharakal]

Fixes three autograd correctness bugs that cause silently frozen or crashing training — the "wrong but no error" class.

#862 — CrossEntropyLoss detaches the tape

The index-target path built its result host-side (tensorDataFactory.init + fromData), reading logProbs values out of the tensor. The resulting loss had no recorded parents, so backward stopped there and the predictions got a null/zero gradient — training froze at constant loss with no error. The soft-target path had a subtler variant: targets * logProbs dispatches through targets.ops, so it only recorded when the targets happened to live on the recording context.

Fix: compute the NLL with differentiable ops dispatched through the predictions' ops. The index path builds a constant one-hot selector (reading discrete indices host-side is fine — they carry no gradient) and computes -sum(oneHot * logProbs); the soft path multiplies via logProbs.ops. Loss values are unchanged; the tape stays connected.

#863 — softmax/logSoftmax backward crashes on negative dim, rank ≥ 3

softmax(dim = -1) passed -1 straight to broadcastToInput, which reserves -1 as a "no-unsqueeze" sentinel — so the reduced axis was never re-expanded and the subsequent elementwise op threw Shapes [2,2,4,4] and [2,2,4] cannot be broadcasted. Any attention-style model using the natural softmax(dim = -1) spelling crashed in backward.

Fix: softmaxGrad/logSoftmaxGrad normalize a negative dim (dim + rank) before use. Softmax always reduces exactly one real axis, so this is unambiguous and leaves the sentinel semantics of sumGrad/meanGrad intact.

#864 — varianceBackward mis-broadcasts on rank ≥ 2

mean(x, dim) dropped the reduced axis (no keepdim), so x - mean could not broadcast; N used the total volume instead of the axis size; and the reduced-shape upstream was never expanded back. A LayerNorm built as (x - mean) / sqrt(variance(x, dim) + eps) crashed in backward.

Fix: re-expand the mean and upstream over the reduced axis (unsqueeze / broadcastToInput) and use the axis size for N.

Tests

• OpsAutodiffBackwardTest: finite-difference backward checks for softmax(dim=-1), logSoftmax(dim=-1), and variance(dim=2) on rank-3 tensors (non-uniform upstream so a wrong gradient is detectable).
• CrossEntropyBackwardTest (new): index-target CE gradient matches the exact analytic (softmax − oneHot)/N; soft-target CE records a non-zero gradient even when the targets live on a separate eager context.
• Full jvmTest for skainet-lang-core, skainet-compile-dag, and skainet-backend-cpu green.

Found while porting an educational GPT trainer to SKaiNET; each bug had a precise reproduction there (training froze / crashed until worked around). This removes the need for those workarounds.

[github-actions]

📖 Documentation Preview

The documentation has been built successfully for this PR.

Generated Files:

• Operator documentation: docs/modules/operators/_generated_/
• JSON schema output: operators.json

Artifacts:

• Download the documentation-preview-877 artifact to view the complete documentation locally.

This comment will be updated automatically when the PR is updated.