diff --git a/skainet-compile/skainet-compile-dag/src/commonMain/kotlin/sk/ainet/lang/graph/DefaultExecutionTape.kt b/skainet-compile/skainet-compile-dag/src/commonMain/kotlin/sk/ainet/lang/graph/DefaultExecutionTape.kt index 3c279d8f5..6ff38f3d9 100644 --- a/skainet-compile/skainet-compile-dag/src/commonMain/kotlin/sk/ainet/lang/graph/DefaultExecutionTape.kt +++ b/skainet-compile/skainet-compile-dag/src/commonMain/kotlin/sk/ainet/lang/graph/DefaultExecutionTape.kt @@ -941,17 +941,22 @@ public class DefaultGradientTape( } override fun varianceBackward(upstream: Tensor, output: Tensor, inputs: List>, attributes: Map): List?> { - // var(x) = E[x^2] - E[x]^2 - // d(var(x))/dx = 2/N * (x - E[x]) + // Biased variance: var(x) = E[(x - E[x])^2]; d(var)/dx = (2/N)(x - E[x]). + // Both the mean and the upstream gradient carry the reduced shape, so + // they are expanded back over the reduced axis before combining — + // otherwise the subtract/multiply cannot broadcast for rank >= 2. val x = inputs[0] - val dim = (attributes["dim"] as? Int) - val meanX = x.ops.mean(x, dim) - // Need to broadcast meanX back to x shape. - // For now, let's assume it's handled or use a simplified version. - val diff = x.ops.subtract(x, meanX) - val n = x.shape.volume.toDouble() // Simplified: should be size along dim + val nd = (attributes["dim"] as? Int)?.let { if (it < 0) it + x.rank else it } + + val meanX = x.ops.mean(x, nd) + val meanKeep = if (nd != null) x.ops.unsqueeze(meanX, nd) else meanX + val diff = x.ops.subtract(x, meanKeep) + + val n = (if (nd != null) x.shape[nd] else x.shape.volume).coerceAtLeast(1) val gradX = diff.ops.mulScalar(diff, 2.0 / n) - return listOf(upstream.ops.multiply(upstream, gradX)) + + val upstreamFull = broadcastToInput(upstream, x, nd) + return listOf(gradX.ops.multiply(gradX, upstreamFull)) } override fun sqrtBackward(upstream: Tensor, output: Tensor, inputs: List>, attributes: Map): List?> { @@ -1454,17 +1459,22 @@ public class DefaultGradientTape( } private fun softmaxGrad(upstream: Tensor, output: Tensor, dim: Int): Tensor { + // softmax reduces exactly one axis: normalize a negative dim (e.g. -1 + // from `softmax(dim = -1)`) so the reduced axis is re-expanded correctly + // by broadcastToInput, which reserves -1 as a "no-unsqueeze" sentinel. + val d = if (dim < 0) dim + output.rank else dim val dot = upstream.ops.multiply(upstream, output) - val sum = upstream.ops.sum(dot, dim) - val expanded = broadcastToInput(sum, output, dim) + val sum = upstream.ops.sum(dot, d) + val expanded = broadcastToInput(sum, output, d) val diff = upstream.ops.subtract(upstream, expanded) return output.ops.multiply(output, diff) } private fun logSoftmaxGrad(upstream: Tensor, logOutput: Tensor, dim: Int): Tensor { + val d = if (dim < 0) dim + logOutput.rank else dim val softmaxApprox = expTensor(logOutput) - val sum = upstream.ops.sum(upstream, dim) - val expanded = broadcastToInput(sum, logOutput, dim) + val sum = upstream.ops.sum(upstream, d) + val expanded = broadcastToInput(sum, logOutput, d) val scaled = softmaxApprox.ops.multiply(softmaxApprox, expanded) return upstream.ops.subtract(upstream, scaled) } diff --git a/skainet-compile/skainet-compile-dag/src/commonTest/kotlin/sk/ainet/exec/autograd/CrossEntropyBackwardTest.kt b/skainet-compile/skainet-compile-dag/src/commonTest/kotlin/sk/ainet/exec/autograd/CrossEntropyBackwardTest.kt new file mode 100644 index 000000000..169e7af5a --- /dev/null +++ b/skainet-compile/skainet-compile-dag/src/commonTest/kotlin/sk/ainet/exec/autograd/CrossEntropyBackwardTest.kt @@ -0,0 +1,109 @@ +package sk.ainet.exec.autograd + +import kotlin.math.abs +import kotlin.math.exp +import kotlin.test.Test +import kotlin.test.assertNotNull +import kotlin.test.assertTrue +import sk.ainet.context.DirectCpuExecutionContext +import sk.ainet.context.Phase +import sk.ainet.exec.tensor.ops.DefaultCpuOps +import sk.ainet.lang.graph.DefaultComputeGraph +import sk.ainet.lang.graph.DefaultGradientTape +import sk.ainet.lang.graph.DefaultGraphExecutionContext +import sk.ainet.lang.nn.loss.CrossEntropyLoss +import sk.ainet.lang.nn.loss.Reduction +import sk.ainet.lang.tensor.Shape +import sk.ainet.lang.tensor.Tensor +import sk.ainet.lang.tensor.data.DenseTensorDataFactory +import sk.ainet.lang.tensor.data.FloatArrayTensorData +import sk.ainet.lang.tensor.withRequiresGrad +import sk.ainet.lang.trace.GraphSink +import sk.ainet.lang.types.FP32 +import sk.ainet.lang.types.Int32 + +/** + * Regression for issue #862: CrossEntropyLoss must backpropagate to the + * predictions. The index-target path used to build its result host-side (via + * tensorDataFactory + fromData), which detached the tape and produced null / + * zero gradients — training silently froze. The soft-target path only recorded + * when the targets lived on the recording context. + */ +class CrossEntropyBackwardTest { + + private fun graphCtx(): DefaultGraphExecutionContext { + val dataFactory = DenseTensorDataFactory() + val graph = DefaultComputeGraph() + return DefaultGraphExecutionContext( + baseOps = DefaultCpuOps(dataFactory), + phase = Phase.TRAIN, + tensorDataFactory = dataFactory, + createTapeFactory = { _ -> DefaultGradientTape(true) }, + computeGraph = graph, + baseSink = GraphSink(graph), + ) + } + + private fun buf(t: Tensor<*, *>): FloatArray = (t.data as FloatArrayTensorData<*>).buffer + + private fun softmaxRow(logits: FloatArray): FloatArray { + val m = logits.max() + val ex = logits.map { exp((it - m).toDouble()) } + val s = ex.sum() + return FloatArray(logits.size) { (ex[it] / s).toFloat() } + } + + @Test + fun index_target_cross_entropy_backprops_to_predictions() { + val c = graphCtx() + val predsFlat = floatArrayOf(1f, 2f, 3f, 1f, 1f, 1f) // [2,3] + val preds = c.fromFloatArray(Shape(2, 3), FP32::class, predsFlat).withRequiresGrad() + val targets = c.fromIntArray(Shape(2), Int32::class, intArrayOf(2, 0)) + + val pair = c.record { + CrossEntropyLoss().forward(preds, targets, this, Reduction.MEAN) + } + val loss = pair.second + (pair.first as DefaultGradientTape).computeGradients(targets = listOf(loss), sources = listOf(preds)) + + val grad = preds.grad + assertNotNull(grad, "CrossEntropyLoss must populate preds.grad (was detached in #862)") + + // Analytic mean-CE gradient: (softmax(row) - oneHot) / N + val n = 2 + val expected = FloatArray(6) + val cls = intArrayOf(2, 0) + for (row in 0 until 2) { + val sm = softmaxRow(predsFlat.copyOfRange(row * 3, row * 3 + 3)) + for (j in 0 until 3) { + expected[row * 3 + j] = (sm[j] - if (j == cls[row]) 1f else 0f) / n + } + } + val actual = buf(grad) + for (i in expected.indices) { + assertTrue(abs(actual[i] - expected[i]) < 1e-4f, "[$i] expected ${expected[i]} got ${actual[i]}") + } + } + + @Test + fun soft_target_cross_entropy_records_even_with_eager_targets() { + val c = graphCtx() + val predsFlat = floatArrayOf(1f, 3f, 2f, 0f) // [2,2] + val preds = c.fromFloatArray(Shape(2, 2), FP32::class, predsFlat).withRequiresGrad() + + // Targets created on a *separate, eager* context — the soft-target multiply + // must still record through the predictions' ops. + val eager = DirectCpuExecutionContext() + val targets = eager.fromFloatArray(Shape(2, 2), FP32::class, floatArrayOf(0.25f, 0.75f, 0.6f, 0.4f)) + + val pair = c.record { + CrossEntropyLoss().forward(preds, targets, this, Reduction.MEAN) + } + (pair.first as DefaultGradientTape).computeGradients(targets = listOf(pair.second), sources = listOf(preds)) + + val grad = preds.grad + assertNotNull(grad, "soft-target CrossEntropyLoss must populate preds.grad") + // Gradient must be non-trivial (not all zeros) + assertTrue(buf(grad).any { abs(it) > 1e-6f }, "gradient should be non-zero") + } +} diff --git a/skainet-compile/skainet-compile-dag/src/commonTest/kotlin/sk/ainet/exec/autograd/OpsAutodiffBackwardTest.kt b/skainet-compile/skainet-compile-dag/src/commonTest/kotlin/sk/ainet/exec/autograd/OpsAutodiffBackwardTest.kt index 578323806..a78599dcb 100644 --- a/skainet-compile/skainet-compile-dag/src/commonTest/kotlin/sk/ainet/exec/autograd/OpsAutodiffBackwardTest.kt +++ b/skainet-compile/skainet-compile-dag/src/commonTest/kotlin/sk/ainet/exec/autograd/OpsAutodiffBackwardTest.kt @@ -174,4 +174,37 @@ class OpsAutodiffBackwardTest { x.ops.convTranspose1d(x, wT, null, stride = 1, padding = 0, outputPadding = 0, dilation = 1, groups = 1) } } + + // ── regression: softmax/logSoftmax backward with a negative dim on rank>=3 (issue #863) ── + + @Test + fun softmax_backward_negative_dim_rank3_matches_finite_diff() { + // `softmax(dim = -1)` on a [2,2,3] tensor previously crashed in backward + // because broadcastToInput skipped re-expanding the reduced (negative) axis. + assertGradMatchesFiniteDiff(Shape(2, 2, 3), FloatArray(12) { (it - 6) * 0.2f }, tol = 1e-2f) { c, x -> + val w = floatTensor(c, Shape(2, 2, 3), FloatArray(12) { 1f + it * 0.1f }) + x.ops.multiply(x.ops.softmax(x, dim = -1), w) // non-uniform upstream so a wrong grad is detectable + } + } + + @Test + fun logSoftmax_backward_negative_dim_rank3_matches_finite_diff() { + assertGradMatchesFiniteDiff(Shape(2, 2, 3), FloatArray(12) { (it - 6) * 0.2f }, tol = 1e-2f) { c, x -> + val w = floatTensor(c, Shape(2, 2, 3), FloatArray(12) { 1f + it * 0.1f }) + x.ops.multiply(x.ops.logSoftmax(x, dim = -1), w) + } + } + + // ── regression: variance backward on rank>=3 with a reduced axis (issue #864) ── + + @Test + fun variance_backward_rank3_matches_finite_diff() { + // `variance(dim = 2)` on a [2,2,3] tensor previously threw "shapes cannot be + // broadcasted" because the reduced mean/upstream were not re-expanded. + assertGradMatchesFiniteDiff(Shape(2, 2, 3), FloatArray(12) { (it - 5) * 0.3f }, tol = 2e-2f) { c, x -> + val v = x.ops.variance(x, dim = 2) // [2,2] + val w = floatTensor(c, Shape(2, 2), floatArrayOf(1f, 0.5f, 1.5f, 0.7f)) + x.ops.multiply(v, w) + } + } } diff --git a/skainet-lang/skainet-lang-core/src/commonMain/kotlin/sk/ainet/lang/nn/loss/CrossEntropyLoss.kt b/skainet-lang/skainet-lang-core/src/commonMain/kotlin/sk/ainet/lang/nn/loss/CrossEntropyLoss.kt index 6d23a0ce4..67d48e459 100644 --- a/skainet-lang/skainet-lang-core/src/commonMain/kotlin/sk/ainet/lang/nn/loss/CrossEntropyLoss.kt +++ b/skainet-lang/skainet-lang-core/src/commonMain/kotlin/sk/ainet/lang/nn/loss/CrossEntropyLoss.kt @@ -56,17 +56,26 @@ public class CrossEntropyLoss @kotlin.jvm.JvmOverloads constructor( validateIndexTargetShapes(preds, targets, classDim) val logProbs = preds.logSoftmax(classDim) - val outData = ctx.tensorDataFactory.init(targets.shape, preds.dtype) { idx -> - val cls = targets.data.get(*idx) as Int + + // Build a one-hot selector for the target classes as a constant tensor, + // then compute the NLL with differentiable ops: -sum_c(oneHot * logProbs). + // Reading the class indices host-side to construct the (non-differentiable) + // one-hot is fine; the gradient path stays intact because the multiply and + // sum are recorded through the predictions' ops — unlike the previous + // host-side result construction, which detached the tape. + val oneHotData = ctx.tensorDataFactory.init(preds.shape, preds.dtype) { idx -> + val sampleIdx = removeClassIndex(idx, classDim) + val cls = targets.data.get(*sampleIdx) as Int require(cls in 0 until classCount) { "CrossEntropyLoss target index $cls out of range [0, $classCount)" } - val logIdx = insertClassIndex(idx, cls, classDim, preds.rank) - val logVal = logProbs.data.get(*logIdx) as Float @Suppress("UNCHECKED_CAST") - (-logVal) as V + (if (idx[classDim] == cls) 1f else 0f) as V } - return ctx.fromData(outData, preds.dtype) + val oneHot = ctx.fromData(oneHotData, preds.dtype) + val weighted = logProbs.ops.multiply(logProbs, oneHot) + val perSample = weighted.ops.sum(weighted, classDim) + return perSample.ops.mulScalar(perSample, -1.0) } private fun computeSoftTargetLoss( @@ -79,29 +88,23 @@ public class CrossEntropyLoss @kotlin.jvm.JvmOverloads constructor( "CrossEntropyLoss with soft targets requires preds/targets shape match, got ${preds.shape.dimensions.contentToString()} vs ${targets.shape.dimensions.contentToString()}" } val logProbs = preds.logSoftmax(classDim) - val weighted = targets * logProbs - val summed = weighted.sum(classDim) - return (-1f) * summed + // Dispatch the multiply through the predictions' ops (not `targets * ...`, + // which would route through targets.ops and detach the tape when the + // targets live on a non-recording context). + val weighted = logProbs.ops.multiply(logProbs, targets) + val summed = weighted.ops.sum(weighted, classDim) + return summed.ops.mulScalar(summed, -1.0) } - private fun insertClassIndex( - baseIdx: IntArray, - cls: Int, - classDim: Int, - outRank: Int - ): IntArray { - val result = IntArray(outRank) - var bi = 0 - for (i in 0 until outRank) { - if (i == classDim) { - result[i] = cls - } else { - result[i] = baseIdx[bi++] - } - } + /** Drops the class-dimension coordinate from a full prediction index. */ + private fun removeClassIndex(fullIdx: IntArray, classDim: Int): IntArray { + val result = IntArray(fullIdx.size - 1) + var j = 0 + for (i in fullIdx.indices) if (i != classDim) result[j++] = fullIdx[i] return result } + private fun validateIndexTargetShapes( preds: Tensor<*, *>, targets: Tensor<*, *>,