Gradient tape: negative dim breaks softmax/logSoftmax backward (broadcastToInput skips unsqueeze)

[michalharakal]

Summary

DefaultGradientTape.broadcastToInput only re-inserts the reduced dimension for non-negative dims:

val targetDim = dim ?: -1
if (targetDim >= 0 && targetDim <= input.rank) {
    expanded = expanded.ops.unsqueeze(expanded, targetDim)
}

softmaxGrad/logSoftmaxGrad pass the op's dim attribute straight through. When the forward pass used the default softmax(dim = -1), the backward computes sum(upstream * y, dim = -1) — which drops the last axis — and then skips the unsqueeze, so the subsequent elementwise op fails:

java.lang.IllegalArgumentException: Shapes [2, 2, 4, 4] and [2, 2, 4] cannot be broadcasted
    at DefaultCpuOpsBase.broadcastShapes
    at DefaultGradientTape.broadcastToInput / softmaxGrad

Any attention-style model that calls softmax(dim = -1) (the natural spelling) crashes in backward on rank ≥ 3 tensors. Workaround: always pass a normalized non-negative dim (softmax(dim = t.rank - 1)).

Suggested fix

Normalize in broadcastToInput (and audit other dim-carrying backward rules): val targetDim = ((dim ?: -1).let { if (it < 0) it + input.rank else it }). Happy to send a PR.

Found on 0.36.0 while backpropagating through a from-scratch multi-head attention.

[michalharakal]

Fixed by the merged PR (#877 for #863/#864, #879 for #859); the auto-close keyword didn't fire because the PR addressed several issues at once.