Summary
Backpropagating through variance(dim) fails for multi-dimensional inputs:
java.lang.IllegalArgumentException: Shapes [2, 4, 8] and [2, 4] cannot be broadcasted
at DefaultCpuOpsBase.subtract
at DefaultGradientTape.varianceBackward
varianceBackward computes x - mean(x, dim) where the reduced mean has lost the reduced axis, so the subtract cannot broadcast against the original input (right-aligned broadcasting misplaces the surviving axes). A LayerNorm implemented as (x - mean) / sqrt(variance(x, dim) + eps) on a [B, T, C] activation crashes in backward.
Reproduce (0.36.0)
Record variance(dim = 2) of a [2, 4, 8] FP32 tensor on a DefaultGraphExecutionContext and call backward — the exception above.
Suggested fix
Unsqueeze the reduced axis before the subtract (keepdim semantics), mirroring what the forward math requires: mean(x, dim).unsqueeze(dim). Same pattern as issue about broadcastToInput and negative dims — several reduction backwards assume rank-1/2 layouts.
Workaround used in our project: compose variance from mean/subtract/multiply ops (E[(x - E[x])^2]), whose backwards are correct. Happy to send a PR.
Summary
Backpropagating through
variance(dim)fails for multi-dimensional inputs:varianceBackwardcomputesx - mean(x, dim)where the reduced mean has lost the reduced axis, so the subtract cannot broadcast against the original input (right-aligned broadcasting misplaces the surviving axes). A LayerNorm implemented as(x - mean) / sqrt(variance(x, dim) + eps)on a[B, T, C]activation crashes in backward.Reproduce (0.36.0)
Record
variance(dim = 2)of a[2, 4, 8]FP32 tensor on aDefaultGraphExecutionContextand callbackward— the exception above.Suggested fix
Unsqueeze the reduced axis before the subtract (keepdim semantics), mirroring what the forward math requires:
mean(x, dim).unsqueeze(dim). Same pattern as issue aboutbroadcastToInputand negative dims — several reduction backwards assume rank-1/2 layouts.Workaround used in our project: compose variance from
mean/subtract/multiplyops (E[(x - E[x])^2]), whose backwards are correct. Happy to send a PR.