diff --git a/CHANGELOG.md b/CHANGELOG.md index dfa071c9..1dce90e9 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -2,6 +2,11 @@ ## [Unreleased] +### Fixed + +- **`concat` backward with inputs of different sizes** (`skainet-compile-dag`, #1490): the rule split the upstream gradient by one input's size and picked chunk `offset / size`, which only lines up when every input has the same extent along the concat axis; other shapes got another input's slice, or `null`, without an error. Each input now receives its own `narrow` of the upstream at the offset the forward gave it, for any sizes and a negative `dim`. +- **`squeeze` backward with no axis** (`skainet-compile-dag`, #1490): `squeeze(x)` removes every size-1 axis, but the backward unsqueezed dim 0 alone, so `x: [1, 3, 1]` got a `[1, 3]` gradient. The gradient is now the upstream reshaped to the input's shape. + ### Added - **`scaledDotProductAttention` backward** (`skainet-compile-dag`, #1384): the attention adjoint (dQ, dK, dV and the dense mask's gradient) replaces the stub that returned `null`, so attention layers train. It honours `scale`, `causal`, sliding `window` / `rightContext`, `keyLengths` and the dense mask, recomputes the probabilities instead of storing them, and reduces grouped-query heads into their shared K/V. Masked scores get exactly zero gradient. `SdpaBackwardTest` checks every option against finite differences and three cases against PyTorch autograd; `stubbedBackwardOpNames` is empty and the coverage matrix's Backward column is complete. 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 8fd1e443..2a02d63c 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 @@ -854,35 +854,19 @@ public class DefaultGradientTape( override fun flattenBackward(upstream: Tensor, output: Tensor, inputs: List>, attributes: Map): List?> = listOf(upstream.ops.reshape(upstream, inputs[0].shape)) override fun concatBackward(upstream: Tensor, output: Tensor, inputs: List>, attributes: Map): List?> { - val dim = (attributes["dim"] as? Int) ?: 0 - val grads = mutableListOf?>() + val rawDim = (attributes["dim"] as? Int) ?: 0 + val dim = if (rawDim < 0) upstream.shape.rank + rawDim else rawDim + // Each input's gradient is its own slice of the upstream along `dim`, at the offset where + // the forward placed it. `narrow` takes any slice length, so inputs of different sizes + // along `dim` get the right slice (splitting the upstream by one input's size and + // indexing `offset / size` only worked when every input had the same size). var offset = 0 - for (input in inputs) { + return inputs.map { input -> val size = input.shape[dim] - if (size == 0) { - grads.add(null) - continue - } - try { - // Since DefaultCpuOps.split(tensor, splitSize, dim) produces chunks of splitSize, - // and the last chunk might be smaller, we can only use it if all chunks except - // possibly the last one match splitSize. - // However, concat might have arbitrary sizes. - // For now, let's try to slice it if possible, but we don't have slice yet. - // If all sizes are equal, split(upstream, size, dim) works perfectly. - val chunks = upstream.ops.split(upstream, size, dim) - val chunkIndex = offset / size - if (chunkIndex < chunks.size) { - grads.add(chunks[chunkIndex]) - } else { - grads.add(null) - } - } catch (e: Exception) { - grads.add(null) - } + val grad = if (size == 0) null else upstream.ops.narrow(upstream, dim, offset, size) offset += size + grad } - return grads } override fun splitBackward(upstream: Tensor, output: Tensor, inputs: List>, attributes: Map): List?> { @@ -892,7 +876,9 @@ public class DefaultGradientTape( // Kept here only to satisfy the DifferentiableTensorOps interface. return listOf(null) } - override fun squeezeBackward(upstream: Tensor, output: Tensor, inputs: List>, attributes: Map): List?> = listOf(upstream.ops.unsqueeze(upstream, (attributes["dim"] as? Int) ?: 0)) // simplistic + // squeeze only drops size-1 axes, so the gradient is the upstream laid out in the input's + // shape — whichever axes `dim = null` removed (unsqueezing dim 0 alone got the shape wrong). + override fun squeezeBackward(upstream: Tensor, output: Tensor, inputs: List>, attributes: Map): List?> = listOf(upstream.ops.reshape(upstream, inputs[0].shape)) override fun unsqueezeBackward(upstream: Tensor, output: Tensor, inputs: List>, attributes: Map): List?> = listOf(upstream.ops.squeeze(upstream, (attributes["dim"] as? Int) ?: 0)) override fun sigmoidBackward(upstream: Tensor, output: Tensor, inputs: List>, attributes: Map): List?> { // d(sigmoid(x))/dx = sigmoid(x) * (1 - sigmoid(x)) = output * (1 - output) diff --git a/skainet-compile/skainet-compile-dag/src/commonTest/kotlin/sk/ainet/exec/autograd/ConcatSqueezeBackwardTest.kt b/skainet-compile/skainet-compile-dag/src/commonTest/kotlin/sk/ainet/exec/autograd/ConcatSqueezeBackwardTest.kt new file mode 100644 index 00000000..d4adc3a7 --- /dev/null +++ b/skainet-compile/skainet-compile-dag/src/commonTest/kotlin/sk/ainet/exec/autograd/ConcatSqueezeBackwardTest.kt @@ -0,0 +1,140 @@ +package sk.ainet.exec.autograd + +import kotlin.test.Test +import kotlin.test.assertContentEquals +import kotlin.test.assertEquals +import kotlin.test.assertNotNull +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.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 + +/** + * `concatBackward` and `squeezeBackward` on the shapes that used to go wrong silently: inputs of + * different sizes along the concat axis, and `squeeze` with no axis. The loss is a weighted sum, + * so each position's gradient is the weight that sat at that position — an exact expectation. + */ +class ConcatSqueezeBackwardTest { + + private fun ctx(): 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 DefaultGraphExecutionContext.floats(shape: Shape, values: FloatArray): Tensor = + fromFloatArray(shape, FP32::class, values) + + private fun buf(t: Tensor<*, *>): FloatArray = (t.data as FloatArrayTensorData<*>).buffer + + /** Gradients of [inputs] for `sum(f(inputs) ⊙ weights)`, where the weights count up from 1 in row-major order. */ + private fun grads( + c: DefaultGraphExecutionContext, + inputs: List>, + f: DefaultGraphExecutionContext.(List>) -> Tensor, + ): List> { + val (tape, loss) = c.record { + val out = f(this as DefaultGraphExecutionContext, inputs) + val w = (this as DefaultGraphExecutionContext).floats(out.shape, FloatArray(out.shape.volume) { (it + 1).toFloat() }) + val weighted = out.ops.multiply(out, w) + weighted.ops.sum(weighted) + } + (tape as DefaultGradientTape).computeGradients(targets = listOf(loss), sources = inputs) + return inputs.mapIndexed { i, x -> assertNotNull(x.grad, "input $i has no gradient") } + } + + private fun assertGrad(expectedShape: Shape, expected: FloatArray, actual: Tensor) { + assertEquals(expectedShape.dimensions.toList(), actual.shape.dimensions.toList(), "gradient shape") + assertContentEquals(expected, buf(actual)) + } + + // ── concat ──────────────────────────────────────────────────────────────────────────────── + + @Test + fun concat_unequal_sizes_along_dim_1() { + val c = ctx() + val a = c.floats(Shape(2, 1), floatArrayOf(1f, 2f)).withRequiresGrad() + val b = c.floats(Shape(2, 2), floatArrayOf(3f, 4f, 5f, 6f)).withRequiresGrad() + // out = [[a0, b0, b1], [a1, b2, b3]] with weights 1..6 + val g = grads(c, listOf(a, b)) { (x, y) -> ops.concat(listOf(x, y), 1) } + assertGrad(Shape(2, 1), floatArrayOf(1f, 4f), g[0]) + assertGrad(Shape(2, 2), floatArrayOf(2f, 3f, 5f, 6f), g[1]) + } + + @Test + fun concat_three_inputs_of_different_sizes_along_dim_0() { + val c = ctx() + val a = c.floats(Shape(1, 2), floatArrayOf(1f, 2f)).withRequiresGrad() + val b = c.floats(Shape(3, 2), FloatArray(6) { it.toFloat() }).withRequiresGrad() + val d = c.floats(Shape(2, 2), FloatArray(4) { it.toFloat() }).withRequiresGrad() + // out rows: a (1), b (3), d (2) → 6 rows, weights 1..12 + val g = grads(c, listOf(a, b, d)) { (x, y, z) -> ops.concat(listOf(x, y, z), 0) } + assertGrad(Shape(1, 2), floatArrayOf(1f, 2f), g[0]) + assertGrad(Shape(3, 2), floatArrayOf(3f, 4f, 5f, 6f, 7f, 8f), g[1]) + assertGrad(Shape(2, 2), floatArrayOf(9f, 10f, 11f, 12f), g[2]) + } + + @Test + fun concat_equal_sizes_still_split_correctly() { + val c = ctx() + val a = c.floats(Shape(2, 2), floatArrayOf(1f, 2f, 3f, 4f)).withRequiresGrad() + val b = c.floats(Shape(2, 2), floatArrayOf(5f, 6f, 7f, 8f)).withRequiresGrad() + val g = grads(c, listOf(a, b)) { (x, y) -> ops.concat(listOf(x, y), 1) } + assertGrad(Shape(2, 2), floatArrayOf(1f, 2f, 5f, 6f), g[0]) + assertGrad(Shape(2, 2), floatArrayOf(3f, 4f, 7f, 8f), g[1]) + } + + @Test + fun concat_negative_dim() { + val c = ctx() + val a = c.floats(Shape(2, 1), floatArrayOf(1f, 2f)).withRequiresGrad() + val b = c.floats(Shape(2, 2), floatArrayOf(3f, 4f, 5f, 6f)).withRequiresGrad() + val g = grads(c, listOf(a, b)) { (x, y) -> ops.concat(listOf(x, y), -1) } + assertGrad(Shape(2, 1), floatArrayOf(1f, 4f), g[0]) + assertGrad(Shape(2, 2), floatArrayOf(2f, 3f, 5f, 6f), g[1]) + } + + // ── squeeze ─────────────────────────────────────────────────────────────────────────────── + + @Test + fun squeeze_all_axes_restores_the_input_shape() { + val c = ctx() + val x = c.floats(Shape(1, 3, 1), floatArrayOf(1f, 2f, 3f)).withRequiresGrad() + val g = grads(c, listOf(x)) { (v) -> ops.squeeze(v) } + assertGrad(Shape(1, 3, 1), floatArrayOf(1f, 2f, 3f), g[0]) + } + + @Test + fun squeeze_inner_axis_restores_the_input_shape() { + val c = ctx() + val x = c.floats(Shape(2, 1, 2), floatArrayOf(1f, 2f, 3f, 4f)).withRequiresGrad() + val g = grads(c, listOf(x)) { (v) -> ops.squeeze(v, 1) } + assertGrad(Shape(2, 1, 2), floatArrayOf(1f, 2f, 3f, 4f), g[0]) + } + + @Test + fun squeeze_then_concat_chain() { + val c = ctx() + val x = c.floats(Shape(1, 2), floatArrayOf(1f, 2f)).withRequiresGrad() + val y = c.floats(Shape(3), floatArrayOf(3f, 4f, 5f)).withRequiresGrad() + // squeeze(x) → [2], concat with y → [5], weights 1..5 + val g = grads(c, listOf(x, y)) { (a, b) -> ops.concat(listOf(ops.squeeze(a), b), 0) } + assertGrad(Shape(1, 2), floatArrayOf(1f, 2f), g[0]) + assertGrad(Shape(3), floatArrayOf(3f, 4f, 5f), g[1]) + } +}