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

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<DType, Any>, output: Tensor<DType, Any>, inputs: List<Tensor<DType, Any>>, attributes: Map<String, Any?>): List<Tensor<DType, Any>?> {
-        // 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<DType, Any>, output: Tensor<DType, Any>, inputs: List<Tensor<DType, Any>>, attributes: Map<String, Any?>): List<Tensor<DType, Any>?> {
@@ -1454,17 +1459,22 @@ public class DefaultGradientTape(
     }
 
     private fun <T : DType, V> softmaxGrad(upstream: Tensor<T, V>, output: Tensor<T, V>, dim: Int): Tensor<T, V> {
+        // 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 <T : DType, V> logSoftmaxGrad(upstream: Tensor<T, V>, logOutput: Tensor<T, V>, dim: Int): Tensor<T, V> {
+        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)
     }

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<FP32, Float>(Shape(2, 3), FP32::class, predsFlat).withRequiresGrad()
+        val targets = c.fromIntArray<Int32, Int>(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<FP32, Float>(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<FP32, Float>(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")
+    }
+}

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)
+        }
+    }
 }

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<T, V>(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<T, V>(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 <T : DType, V> 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<*, *>,