diff --git a/docs/modules/ROOT/pages/reference/operators/generated/index.adoc b/docs/modules/ROOT/pages/reference/operators/generated/index.adoc index dd9000036..40f914442 100644 --- a/docs/modules/ROOT/pages/reference/operators/generated/index.adoc +++ b/docs/modules/ROOT/pages/reference/operators/generated/index.adoc @@ -1,6 +1,6 @@ = AI-NET Operators Reference -Generated from version `0.62.4` on 2026-10-30 +Generated from version `0.62.4` on 2026-10-31 == Operators by Modality diff --git a/docs/modules/ROOT/pages/reference/operators/generated/tensorops.adoc b/docs/modules/ROOT/pages/reference/operators/generated/tensorops.adoc index 4a9eefd0b..08e5cf240 100644 --- a/docs/modules/ROOT/pages/reference/operators/generated/tensorops.adoc +++ b/docs/modules/ROOT/pages/reference/operators/generated/tensorops.adoc @@ -451,12 +451,6 @@ fun convTranspose2d(input:Tensor, weight:Tensor, bias:Tensor, stride:Pair, paddi `Tensor` -=== Notes - -TIP: *all backends*: no backend kernel yet; TransposedConv2d.forward depends on it - -TIP: *all backends*: https://github.com/SKaiNET-developers/SKaiNET/issues/1371[#1371] - == maxPool2d [.darc-none]#โœ– Generated facts only (no human prose)# diff --git a/docs/modules/ROOT/pages/reference/ops-status-matrix.adoc b/docs/modules/ROOT/pages/reference/ops-status-matrix.adoc index 2a434d8bc..e0bf3086c 100644 --- a/docs/modules/ROOT/pages/reference/ops-status-matrix.adoc +++ b/docs/modules/ROOT/pages/reference/ops-status-matrix.adoc @@ -1,7 +1,7 @@ = Operator Coverage Matrix :description: Cross-backend status for every operator function in SKaiNET. -Generated from `operators.json` version `0.62.4` on 2026-10-30. +Generated from `operators.json` version `0.62.4` on 2026-10-31. Rows are `Operator.function` pairs. A gap is marked in source with `@NotImplemented(reason = โ€ฆ, issue = โ€ฆ)` (or `@InProgress`) and shows up here as โŒ (or ๐Ÿšง) linked to its tracking issue; the PR that closes the gap removes the annotation. See xref:contributing/issue-taxonomy.adoc[Issue taxonomy] to pick one up. @@ -34,7 +34,7 @@ Rows are `Operator.function` pairs. A gap is marked in source with `@NotImplemen | `TensorOps.conv2d` | โœ– | โ€” | โœ… | โœ… | โ€” | โœ… | โ€” | `TensorOps.conv3d` | โœ– | โ€” | โŒ | โœ… | โ€” | โœ… | โ€” | `TensorOps.convTranspose1d` | โœ– | โ€” | โŒ | โœ… | โ€” | โœ… | โ€” -| `TensorOps.convTranspose2d` | โœ– | โ€” | โŒ | โ€” | โŒ https://github.com/SKaiNET-developers/SKaiNET/issues/1371[#1371] | โŒ https://github.com/SKaiNET-developers/SKaiNET/issues/1371[#1371] | โŒ https://github.com/SKaiNET-developers/SKaiNET/issues/1371[#1371] +| `TensorOps.convTranspose2d` | โœ– | โ€” | โœ… | โœ… | โ€” | โœ… | โ€” | `TensorOps.maxPool2d` | โœ– | โ€” | โœ… | โœ… | โ€” | โœ… | โ€” | `TensorOps.avgPool2d` | โœ– | โ€” | โœ… | โœ… | โ€” | โœ… | โ€” | `TensorOps.upsample2d` | โœ– | โ€” | โœ… | โœ… | โ€” | โœ… | โ€” @@ -83,7 +83,7 @@ Rows are `Operator.function` pairs. A gap is marked in source with `@NotImplemen | `TensorOps.scaledDotProductAttention` | โœ– | composite ยท decompose | โœ… | โŒ stub | โ€” | โœ… | โ€” | `Similarity.cosineDistance` | โœ– | โ€” | โŒ | โ€” | โœ… | โœ… | โœ… -| *Done* | *1 / 67* | *3 / 67* | *50 / 67* | *57 / 58* | *1 / 67* | *66 / 67* | *1 / 67* +| *Done* | *1 / 67* | *3 / 67* | *51 / 67* | *58 / 59* | *1 / 67* | *67 / 67* | *1 / 67* |=== Per-function detail including notes lives in xref:reference/operators/generated/index.adoc[Operator reference]. @@ -111,7 +111,7 @@ Every concrete `sk.ainet.lang.nn` `Module` subclass. โŒ means the layer's `forw | `Lstm` | `sk.ainet.lang.nn` | โœ… | โ€” | `MaxPool2d` | `sk.ainet.lang.nn` | โœ… | โ€” | `TransposedConv1d` | `sk.ainet.lang.nn` | โœ… | โ€” -| `TransposedConv2d` | `sk.ainet.lang.nn` | โŒ | forward needs TensorOps.convTranspose2d, which no backend implements yet โ€” https://github.com/SKaiNET-developers/SKaiNET/issues/1371[#1371] +| `TransposedConv2d` | `sk.ainet.lang.nn` | โœ… | โ€” | `Upsample2d` | `sk.ainet.lang.nn` | โœ… | โ€” | `ActivationsWrapperModule` | `sk.ainet.lang.nn.activations` | โœ… | โ€” | `ELU` | `sk.ainet.lang.nn.activations` | โœ… | โ€” @@ -129,5 +129,5 @@ Every concrete `sk.ainet.lang.nn` `Module` subclass. โŒ means the layer's `forw | `LayerNormalization` | `sk.ainet.lang.nn.normalization` | โœ… | โ€” | `MLP` | `sk.ainet.lang.nn.topology` | โœ… | โ€” -| *Done* | | *31 / 32* | +| *Done* | | *32 / 32* | |=== diff --git a/skainet-backends/skainet-backend-cpu/src/commonMain/kotlin/sk/ainet/exec/tensor/ops/DefaultCpuOps.kt b/skainet-backends/skainet-backend-cpu/src/commonMain/kotlin/sk/ainet/exec/tensor/ops/DefaultCpuOps.kt index 305572c1e..69cd20ef6 100644 --- a/skainet-backends/skainet-backend-cpu/src/commonMain/kotlin/sk/ainet/exec/tensor/ops/DefaultCpuOps.kt +++ b/skainet-backends/skainet-backend-cpu/src/commonMain/kotlin/sk/ainet/exec/tensor/ops/DefaultCpuOps.kt @@ -1814,6 +1814,129 @@ public open class DefaultCpuOpsBase( return newTensor(outData, input.dtype, input) } + override fun convTranspose2d( + input: Tensor, + weight: Tensor, + bias: Tensor?, + stride: Pair, + padding: Pair, + outputPadding: Pair, + dilation: Pair, + groups: Int + ): Tensor { + // input: [batch, C_in, H, W]; weight: [C_in, C_out/groups, kH, kW] (PyTorch ConvTranspose2d layout). + require(input.rank == 4) { "convTranspose2d: input must be 4D (N, C_in, H, W), got ${input.shape.dimensions.contentToString()}" } + require(weight.rank == 4) { "convTranspose2d: weight must be 4D (C_in, C_out/groups, kH, kW), got ${weight.shape.dimensions.contentToString()}" } + require(groups >= 1) { "convTranspose2d: groups must be >= 1" } + require(input.dtype == weight.dtype) { "convTranspose2d: dtype mismatch between input and weight: ${input.dtype} vs ${weight.dtype}" } + bias?.let { require(it.dtype == input.dtype) { "convTranspose2d: dtype mismatch for bias" } } + require(input.dtype == sk.ainet.lang.types.FP32::class || input.dtype == sk.ainet.lang.types.FP16::class) { + "Unsupported dtype for convTranspose2d: ${input.dtype}" + } + + val batch = input.shape[0] + val inChannels = input.shape[1] + val inH = input.shape[2] + val inW = input.shape[3] + + require(weight.shape[0] == inChannels) { "convTranspose2d: weight C_in ${weight.shape[0]} must equal input channels $inChannels" } + require(inChannels % groups == 0) { "convTranspose2d: input channels $inChannels not divisible by groups $groups" } + + val outChannelsPerGroup = weight.shape[1] + val kH = weight.shape[2] + val kW = weight.shape[3] + val outChannels = outChannelsPerGroup * groups + + val (sH, sW) = stride + val (pH, pW) = padding + val (opH, opW) = outputPadding + val (dH, dW) = dilation + require(sH > 0 && sW > 0) { "convTranspose2d: stride must be positive, got ($sH, $sW)" } + require(dH > 0 && dW > 0) { "convTranspose2d: dilation must be positive, got ($dH, $dW)" } + require(pH >= 0 && pW >= 0) { "convTranspose2d: padding must be non-negative, got ($pH, $pW)" } + // PyTorch: output padding must be smaller than either stride or dilation on each axis. + require(opH >= 0 && opH < maxOf(sH, dH) && opW >= 0 && opW < maxOf(sW, dW)) { + "convTranspose2d: outputPadding ($opH, $opW) must be smaller than stride or dilation" + } + + val outH = (inH - 1) * sH - 2 * pH + dH * (kH - 1) + opH + 1 + val outW = (inW - 1) * sW - 2 * pW + dW * (kW - 1) + opW + 1 + require(outH > 0 && outW > 0) { "convTranspose2d: computed non-positive output shape (H=$outH, W=$outW)" } + + if (bias != null) { + when (bias.rank) { + 1 -> require(bias.shape[0] == outChannels) { "convTranspose2d: bias shape must be [C_out], got ${bias.shape.dimensions.contentToString()}" } + 4 -> require(bias.shape[0] == 1 && bias.shape[1] == outChannels && bias.shape[2] == 1 && bias.shape[3] == 1) { + "convTranspose2d: bias shape must be [1,C_out,1,1] when 4D, got ${bias.shape.dimensions.contentToString()}" + } + else -> error("convTranspose2d: unsupported bias rank ${bias.rank}") + } + } + + val outData = dataFactory.zeros(Shape(batch, outChannels, outH, outW), input.dtype) + val inData = input.data + val wData = weight.data + val inChPerGroup = inChannels / groups + + // Scatter-add over the output, mirroring convTranspose1d: + // out[b, oc, ih*sH - pH + kh*dH, iw*sW - pW + kw*dW] += in[b, ic, ih, iw] * w[ic, oc, kh, kw] + for (b in 0 until batch) { + for (g in 0 until groups) { + for (ic in 0 until inChPerGroup) { + val inCh = g * inChPerGroup + ic + for (oc in 0 until outChannelsPerGroup) { + val outCh = g * outChannelsPerGroup + oc + for (ih in 0 until inH) { + val ohBase = ih * sH - pH + for (iw in 0 until inW) { + val inputVal = inData.get(b, inCh, ih, iw) as Float + if (inputVal == 0f) continue + val owBase = iw * sW - pW + for (kh in 0 until kH) { + val oh = ohBase + kh * dH + if (oh < 0 || oh >= outH) continue + for (kw in 0 until kW) { + val ow = owBase + kw * dW + if (ow < 0 || ow >= outW) continue + val weightVal = wData.get(inCh, oc, kh, kw) as Float + val existing = outData.get(b, outCh, oh, ow) as Float + @Suppress("UNCHECKED_CAST") + outData.set(b, outCh, oh, ow, value = (existing + inputVal * weightVal) as V) + } + } + } + } + } + } + } + } + + if (bias != null) { + val biasData = bias.data + for (b in 0 until batch) { + for (oc in 0 until outChannels) { + val biasVal = when (bias.rank) { + 1 -> biasData.get(oc) as Float + else -> biasData.get(0, oc, 0, 0) as Float + } + for (oh in 0 until outH) { + for (ow in 0 until outW) { + val existing = outData.get(b, oc, oh, ow) as Float + @Suppress("UNCHECKED_CAST") + outData.set(b, oc, oh, ow, value = (existing + biasVal) as V) + } + } + } + } + } + + return if (bias != null) { + newTensor(outData, input.dtype, input, weight, bias) + } else { + newTensor(outData, input.dtype, input, weight) + } + } + @TensorOp() override fun conv3d( input: Tensor, diff --git a/skainet-backends/skainet-backend-cpu/src/commonTest/kotlin/sk/ainet/exec/ops/ConvTranspose2dOpsTest.kt b/skainet-backends/skainet-backend-cpu/src/commonTest/kotlin/sk/ainet/exec/ops/ConvTranspose2dOpsTest.kt new file mode 100644 index 000000000..011b7556d --- /dev/null +++ b/skainet-backends/skainet-backend-cpu/src/commonTest/kotlin/sk/ainet/exec/ops/ConvTranspose2dOpsTest.kt @@ -0,0 +1,293 @@ +package sk.ainet.exec.ops + +import kotlin.math.abs +import kotlin.test.Test +import kotlin.test.assertEquals +import kotlin.test.assertFailsWith +import kotlin.test.assertTrue +import sk.ainet.context.DirectCpuExecutionContext +import sk.ainet.context.Phase +import sk.ainet.lang.nn.TransposedConv2d +import sk.ainet.lang.tensor.Shape +import sk.ainet.lang.tensor.Tensor +import sk.ainet.lang.types.FP32 + +class ConvTranspose2dOpsTest { + + private val ctx = DirectCpuExecutionContext(phase = Phase.EVAL) + + private fun tensor(shape: Shape, data: FloatArray): Tensor = + ctx.fromFloatArray(shape, FP32::class, data) + + /** + * Reference conv_transpose2d with torch.nn.functional.conv_transpose2d semantics, + * accumulated in Double. Input is flat NCHW, weight is flat [C_in, C_out/groups, kH, kW]. + * out[b, oc, ih*sH - pH + kh*dH, iw*sW - pW + kw*dW] += in[b, ic, ih, iw] * w[ic, oc, kh, kw] + */ + private fun referenceConvTranspose2d( + input: FloatArray, + n: Int, + cIn: Int, + inH: Int, + inW: Int, + weight: FloatArray, + cOutPerGroup: Int, + kH: Int, + kW: Int, + bias: FloatArray? = null, + stride: Pair = 1 to 1, + padding: Pair = 0 to 0, + outputPadding: Pair = 0 to 0, + dilation: Pair = 1 to 1, + groups: Int = 1, + ): Pair { + val (sH, sW) = stride + val (pH, pW) = padding + val (opH, opW) = outputPadding + val (dH, dW) = dilation + val outH = (inH - 1) * sH - 2 * pH + dH * (kH - 1) + opH + 1 + val outW = (inW - 1) * sW - 2 * pW + dW * (kW - 1) + opW + 1 + val cOut = cOutPerGroup * groups + val inChPerGroup = cIn / groups + val out = DoubleArray(n * cOut * outH * outW) + for (b in 0 until n) { + for (g in 0 until groups) { + for (ic in 0 until inChPerGroup) { + val inCh = g * inChPerGroup + ic + for (oc in 0 until cOutPerGroup) { + val outCh = g * cOutPerGroup + oc + for (ih in 0 until inH) { + for (iw in 0 until inW) { + val x = input[((b * cIn + inCh) * inH + ih) * inW + iw].toDouble() + for (kh in 0 until kH) { + val oh = ih * sH - pH + kh * dH + if (oh < 0 || oh >= outH) continue + for (kw in 0 until kW) { + val ow = iw * sW - pW + kw * dW + if (ow < 0 || ow >= outW) continue + val w = weight[((inCh * cOutPerGroup + oc) * kH + kh) * kW + kw].toDouble() + out[((b * cOut + outCh) * outH + oh) * outW + ow] += x * w + } + } + } + } + } + } + } + } + if (bias != null) { + for (b in 0 until n) { + for (oc in 0 until cOut) { + val base = (b * cOut + oc) * outH * outW + for (i in 0 until outH * outW) out[base + i] += bias[oc].toDouble() + } + } + } + return FloatArray(out.size) { out[it].toFloat() } to Shape(n, cOut, outH, outW) + } + + private fun assertFlatClose(expected: FloatArray, actual: Tensor, shape: Shape, eps: Float = 1e-4f) { + assertEquals(shape, actual.shape, "shape mismatch") + val dims = shape.dimensions + val idx = IntArray(dims.size) + for (flat in expected.indices) { + var rem = flat + for (d in dims.indices.reversed()) { + idx[d] = rem % dims[d] + rem /= dims[d] + } + val a = actual.data.get(*idx) as Float + assertTrue(abs(expected[flat] - a) <= eps, "Mismatch at flat index $flat: expected ${expected[flat]}, got $a") + } + } + + // ========== Hand-verified basics ========== + + @Test + fun convTranspose2d_basic_k2_s1_matchesHandComputation() { + // input [[1,2],[3,4]], ones kernel 2x2, stride 1, no padding -> full scatter sum: + // [[1,3,2],[4,10,6],[3,7,4]] + val input = tensor(Shape(1, 1, 2, 2), floatArrayOf(1f, 2f, 3f, 4f)) + val weight = tensor(Shape(1, 1, 2, 2), floatArrayOf(1f, 1f, 1f, 1f)) + + val output = input.ops.convTranspose2d(input, weight) + + assertFlatClose( + floatArrayOf(1f, 3f, 2f, 4f, 10f, 6f, 3f, 7f, 4f), + output, + Shape(1, 1, 3, 3) + ) + } + + @Test + fun convTranspose2d_stride2_k2_stampsNonOverlappingKernels() { + // k == stride: each input pixel stamps the kernel scaled by its value, no overlap. + val input = tensor(Shape(1, 1, 2, 2), floatArrayOf(1f, 2f, 3f, 4f)) + val weight = tensor(Shape(1, 1, 2, 2), floatArrayOf(1f, 2f, 3f, 4f)) + + val output = input.ops.convTranspose2d(input, weight, stride = 2 to 2) + + assertFlatClose( + floatArrayOf( + 1f, 2f, 2f, 4f, + 3f, 4f, 6f, 8f, + 3f, 6f, 4f, 8f, + 9f, 12f, 12f, 16f, + ), + output, + Shape(1, 1, 4, 4) + ) + } + + // ========== Parameter grid vs double-precision reference ========== + + @Test + fun convTranspose2d_multiChannel_stride2_padding1_outputPadding1_matchesReference() { + val n = 1; val cIn = 2; val inH = 3; val inW = 3 + val cOut = 3; val kH = 3; val kW = 3 + val inputArr = FloatArray(n * cIn * inH * inW) { i -> ((i * 7) % 11) - 5f + i * 0.125f } + val weightArr = FloatArray(cIn * cOut * kH * kW) { i -> ((i * 5) % 13) - 6f + i * 0.0625f } + val input = tensor(Shape(n, cIn, inH, inW), inputArr) + val weight = tensor(Shape(cIn, cOut, kH, kW), weightArr) + + val output = input.ops.convTranspose2d( + input, weight, + stride = 2 to 2, padding = 1 to 1, outputPadding = 1 to 1, + ) + + val (expected, shape) = referenceConvTranspose2d( + inputArr, n, cIn, inH, inW, weightArr, cOut, kH, kW, + stride = 2 to 2, padding = 1 to 1, outputPadding = 1 to 1, + ) + assertFlatClose(expected, output, shape, eps = 1e-3f) + } + + @Test + fun convTranspose2d_asymmetricStrideAndPadding_matchesReference() { + val n = 2; val cIn = 1; val inH = 2; val inW = 4 + val cOut = 2; val kH = 3; val kW = 2 + val inputArr = FloatArray(n * cIn * inH * inW) { i -> (i % 5) - 2f + i * 0.25f } + val weightArr = FloatArray(cIn * cOut * kH * kW) { i -> ((i * 3) % 7) - 3f + i * 0.5f } + val input = tensor(Shape(n, cIn, inH, inW), inputArr) + val weight = tensor(Shape(cIn, cOut, kH, kW), weightArr) + + val output = input.ops.convTranspose2d( + input, weight, + stride = 2 to 3, padding = 1 to 0, outputPadding = 1 to 2, + ) + + val (expected, shape) = referenceConvTranspose2d( + inputArr, n, cIn, inH, inW, weightArr, cOut, kH, kW, + stride = 2 to 3, padding = 1 to 0, outputPadding = 1 to 2, + ) + assertFlatClose(expected, output, shape, eps = 1e-3f) + } + + @Test + fun convTranspose2d_dilation2_matchesReference() { + val n = 1; val cIn = 1; val inH = 3; val inW = 3 + val cOut = 1; val kH = 2; val kW = 2 + val inputArr = FloatArray(9) { i -> i + 1f } + val weightArr = floatArrayOf(1f, -1f, 0.5f, 2f) + val input = tensor(Shape(n, cIn, inH, inW), inputArr) + val weight = tensor(Shape(cIn, cOut, kH, kW), weightArr) + + val output = input.ops.convTranspose2d(input, weight, dilation = 2 to 2) + + val (expected, shape) = referenceConvTranspose2d( + inputArr, n, cIn, inH, inW, weightArr, cOut, kH, kW, + dilation = 2 to 2, + ) + assertFlatClose(expected, output, shape) + } + + @Test + fun convTranspose2d_groups2_matchesReference() { + val n = 1; val cIn = 4; val inH = 2; val inW = 2 + val cOutPerGroup = 2; val kH = 2; val kW = 2; val groups = 2 + val inputArr = FloatArray(n * cIn * inH * inW) { i -> ((i * 3) % 5) - 2f + i * 0.1f } + val weightArr = FloatArray(cIn * cOutPerGroup * kH * kW) { i -> ((i * 7) % 9) - 4f + i * 0.2f } + val input = tensor(Shape(n, cIn, inH, inW), inputArr) + val weight = tensor(Shape(cIn, cOutPerGroup, kH, kW), weightArr) + + val output = input.ops.convTranspose2d(input, weight, stride = 2 to 2, groups = groups) + + val (expected, shape) = referenceConvTranspose2d( + inputArr, n, cIn, inH, inW, weightArr, cOutPerGroup, kH, kW, + stride = 2 to 2, groups = groups, + ) + assertFlatClose(expected, output, shape, eps = 1e-3f) + } + + @Test + fun convTranspose2d_bias_addsPerOutputChannel() { + val input = tensor(Shape(1, 1, 2, 2), floatArrayOf(1f, 2f, 3f, 4f)) + val weight = tensor(Shape(1, 2, 2, 2), FloatArray(8) { 1f }) + val biasArr = floatArrayOf(10f, -10f) + val bias = tensor(Shape(2), biasArr) + + val output = input.ops.convTranspose2d(input, weight, bias) + + val (expected, shape) = referenceConvTranspose2d( + floatArrayOf(1f, 2f, 3f, 4f), 1, 1, 2, 2, FloatArray(8) { 1f }, 2, 2, 2, + bias = biasArr, + ) + assertFlatClose(expected, output, shape) + } + + // ========== Layer integration ========== + + @Test + fun transposedConv2dLayer_forward_matchesOutputSizeAndReference() { + val inCh = 2; val outCh = 1; val k = 4 to 4 + val weightArr = FloatArray(inCh * outCh * 16) { i -> (i % 8) * 0.25f - 1f } + val biasArr = floatArrayOf(0.5f) + val layer = TransposedConv2d( + inChannels = inCh, + outChannels = outCh, + kernelSize = k, + stride = 2 to 2, + padding = 1 to 1, + name = "upconv", + initWeights = tensor(Shape(inCh, outCh, 4, 4), weightArr), + initBias = tensor(Shape(outCh), biasArr), + ) + val inputArr = FloatArray(2 * 3 * 3) { i -> ((i * 5) % 7) - 3f } + val input = tensor(Shape(1, inCh, 3, 3), inputArr) + + val output = layer.forward(input, ctx) + + val (expH, expW) = layer.outputSize(3 to 3) + assertEquals(listOf(1, outCh, expH, expW), output.shape.dimensions.toList()) + val (expected, shape) = referenceConvTranspose2d( + inputArr, 1, inCh, 3, 3, weightArr, outCh, 4, 4, + bias = biasArr, stride = 2 to 2, padding = 1 to 1, + ) + assertFlatClose(expected, output, shape, eps = 1e-3f) + } + + // ========== Validation ========== + + @Test + fun convTranspose2d_wrongWeightChannels_throws() { + val input = tensor(Shape(1, 2, 2, 2), FloatArray(8) { 1f }) + val weight = tensor(Shape(3, 1, 2, 2), FloatArray(12) { 1f }) + assertFailsWith { input.ops.convTranspose2d(input, weight) } + } + + @Test + fun convTranspose2d_outputPaddingNotLessThanStrideAndDilation_throws() { + val input = tensor(Shape(1, 1, 2, 2), FloatArray(4) { 1f }) + val weight = tensor(Shape(1, 1, 2, 2), FloatArray(4) { 1f }) + assertFailsWith { + input.ops.convTranspose2d(input, weight, stride = 1 to 1, outputPadding = 1 to 1) + } + } + + @Test + fun convTranspose2d_non4dInput_throws() { + val input = tensor(Shape(1, 1, 4), FloatArray(4) { 1f }) + val weight = tensor(Shape(1, 1, 2, 2), FloatArray(4) { 1f }) + assertFailsWith { input.ops.convTranspose2d(input, weight) } + } +} diff --git a/skainet-compile/skainet-compile-core/src/commonMain/kotlin/sk/ainet/tape/RecordingExecution.kt b/skainet-compile/skainet-compile-core/src/commonMain/kotlin/sk/ainet/tape/RecordingExecution.kt index a735ae5b5..6182216c4 100644 --- a/skainet-compile/skainet-compile-core/src/commonMain/kotlin/sk/ainet/tape/RecordingExecution.kt +++ b/skainet-compile/skainet-compile-core/src/commonMain/kotlin/sk/ainet/tape/RecordingExecution.kt @@ -484,6 +484,10 @@ internal class RecordingTensorOpsDecorator(private val base: TensorOps) : Tensor input: Tensor, weight: Tensor, bias: Tensor?, stride: Int, padding: Int, outputPadding: Int, dilation: Int, groups: Int ): Tensor = base.convTranspose1d(input, weight, bias, stride, padding, outputPadding, dilation, groups) + override fun convTranspose2d( + input: Tensor, weight: Tensor, bias: Tensor?, + stride: Pair, padding: Pair, outputPadding: Pair, dilation: Pair, groups: Int + ): Tensor = base.convTranspose2d(input, weight, bias, stride, padding, outputPadding, dilation, groups) override fun scaledDotProductAttention( query: Tensor, key: Tensor, value: Tensor, mask: Tensor?, scale: Float, causal: Boolean, diff --git a/skainet-compile/skainet-compile-dag/api/jvm/skainet-compile-dag.api b/skainet-compile/skainet-compile-dag/api/jvm/skainet-compile-dag.api index 6c4474ebd..0cee8f6a2 100644 --- a/skainet-compile/skainet-compile-dag/api/jvm/skainet-compile-dag.api +++ b/skainet-compile/skainet-compile-dag/api/jvm/skainet-compile-dag.api @@ -105,6 +105,7 @@ public final class sk/ainet/lang/graph/DefaultGradientTape : sk/ainet/lang/graph public fun conv2dBackward (Lsk/ainet/lang/tensor/Tensor;Lsk/ainet/lang/tensor/Tensor;Ljava/util/List;Ljava/util/Map;)Ljava/util/List; public fun conv3dBackward (Lsk/ainet/lang/tensor/Tensor;Lsk/ainet/lang/tensor/Tensor;Ljava/util/List;Ljava/util/Map;)Ljava/util/List; public fun convTranspose1dBackward (Lsk/ainet/lang/tensor/Tensor;Lsk/ainet/lang/tensor/Tensor;Ljava/util/List;Ljava/util/Map;)Ljava/util/List; + public fun convTranspose2dBackward (Lsk/ainet/lang/tensor/Tensor;Lsk/ainet/lang/tensor/Tensor;Ljava/util/List;Ljava/util/Map;)Ljava/util/List; public fun copy ()Lsk/ainet/tape/ExecutionTape; public fun cosBackward (Lsk/ainet/lang/tensor/Tensor;Lsk/ainet/lang/tensor/Tensor;Ljava/util/List;Ljava/util/Map;)Ljava/util/List; public fun divScalarBackward (Lsk/ainet/lang/tensor/Tensor;Lsk/ainet/lang/tensor/Tensor;Ljava/util/List;Ljava/util/Map;)Ljava/util/List; 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 358536066..5fbcce466 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 @@ -1237,6 +1237,89 @@ public class DefaultGradientTape( return grads } + override fun convTranspose2dBackward(upstream: Tensor, output: Tensor, inputs: List>, attributes: Map): List?> { + // Adjoint of convTranspose2d forward: out[b,oc,oh,ow] += in[b,ic,ih,iw]*w[ic,oc,kh,kw], + // oh = ih*sH - pH + kh*dH, ow = iw*sW - pW + kw*dW. + val input = inputs[0] + val weight = inputs[1] + val hasBias = inputs.size > 2 + val (sH, sW) = pair2(attributes["stride"], 1) + val (pH, pW) = pair2(attributes["padding"], 0) + val (dH, dW) = pair2(attributes["dilation"], 1) + val groups = (attributes["groups"] as? Number)?.toInt() ?: 1 + + val batch = input.shape[0] + val inChannels = input.shape[1] + val inH = input.shape[2] + val inW = input.shape[3] + val outChannelsPerGroup = weight.shape[1] + val kH = weight.shape[2] + val kW = weight.shape[3] + val outH = upstream.shape[2] + val outW = upstream.shape[3] + val inChPerGroup = inChannels / groups + + val gradInput = zerosLike(input) + val gradWeight = zerosLike(weight) + val gIdx = IntArray(4) + val wIdx = IntArray(4) + val uIdx = IntArray(4) + + for (b in 0 until batch) { + for (g in 0 until groups) { + for (ic in 0 until inChPerGroup) { + val inCh = g * inChPerGroup + ic + for (oc in 0 until outChannelsPerGroup) { + val outCh = g * outChannelsPerGroup + oc + for (ih in 0 until inH) { + for (iw in 0 until inW) { + for (kh in 0 until kH) { + val oh = ih * sH - pH + kh * dH + if (oh < 0 || oh >= outH) continue + for (kw in 0 until kW) { + val ow = iw * sW - pW + kw * dW + if (ow < 0 || ow >= outW) continue + uIdx[0] = b; uIdx[1] = outCh; uIdx[2] = oh; uIdx[3] = ow + val up = (upstream.data.get(*uIdx) as Number).toFloat() + wIdx[0] = inCh; wIdx[1] = oc; wIdx[2] = kh; wIdx[3] = kw + val w = (weight.data.get(*wIdx) as Number).toFloat() + gIdx[0] = b; gIdx[1] = inCh; gIdx[2] = ih; gIdx[3] = iw + val xVal = (input.data.get(*gIdx) as Number).toFloat() + accumulateScalar(gradInput, gIdx, up * w) + accumulateScalar(gradWeight, wIdx, xVal * up) + } + } + } + } + } + } + } + } + + val grads = mutableListOf?>(gradInput, gradWeight) + if (hasBias) { + val bias = inputs[2] + val gradBias = zerosLike(bias) + val bIdx = IntArray(1) + val outChannels = outChannelsPerGroup * groups + for (oc in 0 until outChannels) { + var acc = 0f + for (b in 0 until batch) { + for (oh in 0 until outH) { + for (ow in 0 until outW) { + uIdx[0] = b; uIdx[1] = oc; uIdx[2] = oh; uIdx[3] = ow + acc += (upstream.data.get(*uIdx) as Number).toFloat() + } + } + } + bIdx[0] = oc + accumulateScalar(gradBias, bIdx, acc) + } + grads.add(gradBias) + } + return grads + } + /** Read-add-write [src]'s value at [srcIdx] into [dest] at [destIdx]; used by scatter-add backwards. */ private fun accumulateAt(dest: Tensor, destIdx: IntArray, src: Tensor, srcIdx: IntArray) { accumulateScalar(dest, destIdx, (src.data.get(*srcIdx) as Number).toFloat()) @@ -1302,6 +1385,7 @@ public class DefaultGradientTape( "conv2d" to ::conv2dBackward, "conv3d" to ::conv3dBackward, "convTranspose1d" to ::convTranspose1dBackward, + "convTranspose2d" to ::convTranspose2dBackward, "maxPool2d" to ::maxPool2dBackward, "avgPool2d" to ::avgPool2dBackward, "upsample2d" to ::upsample2dBackward, diff --git a/skainet-compile/skainet-compile-dag/src/commonTest/kotlin/sk/ainet/compile/graph/ComputeGraphExecutorTest.kt b/skainet-compile/skainet-compile-dag/src/commonTest/kotlin/sk/ainet/compile/graph/ComputeGraphExecutorTest.kt index f1d07abd7..ef12c0402 100644 --- a/skainet-compile/skainet-compile-dag/src/commonTest/kotlin/sk/ainet/compile/graph/ComputeGraphExecutorTest.kt +++ b/skainet-compile/skainet-compile-dag/src/commonTest/kotlin/sk/ainet/compile/graph/ComputeGraphExecutorTest.kt @@ -268,6 +268,7 @@ private class TestTensorOps : TensorOps { override fun sin(tensor: Tensor): Tensor = tensor override fun cos(tensor: Tensor): Tensor = tensor override fun convTranspose1d(input: Tensor, weight: Tensor, bias: Tensor?, stride: Int, padding: Int, outputPadding: Int, dilation: Int, groups: Int): Tensor = input + override fun convTranspose2d(input: Tensor, weight: Tensor, bias: Tensor?, stride: Pair, padding: Pair, outputPadding: Pair, dilation: Pair, groups: Int): Tensor = input override fun scaledDotProductAttention(query: Tensor, key: Tensor, value: Tensor, mask: Tensor?, scale: Float, causal: Boolean, keyLengths: Tensor?, window: Int, rightContext: Int): Tensor = query override fun conv1d(input: Tensor, weight: Tensor, bias: Tensor?, stride: Int, padding: Int, dilation: Int, groups: Int): Tensor = input override fun conv2d(input: Tensor, weight: Tensor, bias: Tensor?, stride: Pair, padding: Pair, dilation: Pair, groups: Int): Tensor = input 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 9e6635199..5b6aea38a 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 @@ -218,6 +218,32 @@ class OpsAutodiffBackwardTest { } } + @Test + fun convTranspose2d_backward_matches_finite_diff() { + // input [1,1,2,2], weight [in=1, outPerGroup=1, 2x2], stride 1, padding 0. + val w = floatArrayOf(0.5f, -1.2f, 0.8f, 0.3f) + assertGradMatchesFiniteDiff(Shape(1, 1, 2, 2), floatArrayOf(0.3f, -0.2f, 0.7f, 0.4f)) { c, x -> + val wT = floatTensor(c, Shape(1, 1, 2, 2), w) + x.ops.convTranspose2d( + x, wT, null, + stride = 1 to 1, padding = 0 to 0, outputPadding = 0 to 0, dilation = 1 to 1, groups = 1, + ) + } + } + + @Test + fun convTranspose2d_backward_stride2_padding1_matches_finite_diff() { + // Strided + padded variant: exercises the oh/ow bounds checks in the adjoint. + val w = floatArrayOf(0.5f, -1.2f, 0.8f, 0.3f, -0.4f, 0.9f, 0.2f, -0.6f, 1.1f) + assertGradMatchesFiniteDiff(Shape(1, 1, 2, 2), floatArrayOf(0.3f, -0.2f, 0.7f, 0.4f)) { c, x -> + val wT = floatTensor(c, Shape(1, 1, 3, 3), w) + x.ops.convTranspose2d( + x, wT, null, + stride = 2 to 2, padding = 1 to 1, outputPadding = 1 to 1, dilation = 1 to 1, groups = 1, + ) + } + } + // โ”€โ”€ regression: softmax/logSoftmax backward with a negative dim on rank>=3 (issue #863) โ”€โ”€ @Test diff --git a/skainet-compile/skainet-compile-hlo/src/commonMain/kotlin/sk/ainet/compile/hlo/converters/NeuralNetOperationsConverter.kt b/skainet-compile/skainet-compile-hlo/src/commonMain/kotlin/sk/ainet/compile/hlo/converters/NeuralNetOperationsConverter.kt index 1de313083..5321711c6 100644 --- a/skainet-compile/skainet-compile-hlo/src/commonMain/kotlin/sk/ainet/compile/hlo/converters/NeuralNetOperationsConverter.kt +++ b/skainet-compile/skainet-compile-hlo/src/commonMain/kotlin/sk/ainet/compile/hlo/converters/NeuralNetOperationsConverter.kt @@ -23,7 +23,7 @@ public class NeuralNetOperationsConverter : StableHloOperationConverter { override val supportedOperations: Set = setOf( // Convolutional operations - "conv1d", "conv2d", + "conv1d", "conv2d", "convTranspose2d", // Pooling operations "maxPool2d", "avgPool2d", "averagePool2d", // Normalization operations @@ -45,6 +45,7 @@ public class NeuralNetOperationsConverter : StableHloOperationConverter { return when (node.operation.name.lowercase()) { "conv1d" -> convertConv1d(node, operands, context) "conv2d" -> convertConv2d(node, operands, context) + "convtranspose2d" -> convertConvTranspose2d(node, operands, context) "maxpool2d" -> convertMaxPool2d(node, operands, context) "avgpool2d", "averagepool2d" -> convertAvgPool2d(node, operands, context) "batchnorm", "batchnormalization" -> convertBatchNorm(node, operands, context) @@ -167,6 +168,73 @@ public class NeuralNetOperationsConverter : StableHloOperationConverter { ) } + private fun convertConvTranspose2d( + node: GraphNode, + operands: List, + context: ConversionContext + ): ConversionResult { + if (operands.size < 2) { + return ConversionResult.Failure( + "ConvTranspose2d operation requires at least 2 operands (input, weight), got ${operands.size}", + "Unsupported convTranspose2d arity for node ${node.id}" + ) + } + + val params = node.operation.parameters + val stride = extractStride(params) + val padding = extractPadding(params) + val outputPadding = extractOutputPadding(params) + val dilation = extractDilation(params) + val groups = params["groups"] as? Int ?: 1 + if (groups != 1) { + return ConversionResult.Failure( + "convTranspose2d lowering supports groups == 1 only, got $groups", + "Grouped transposed convolution needs a kernel reshape that is not implemented yet (see #1375)" + ) + } + + // kH/kW come from the weight spec [C_in, C_out, kH, kW]; the transposed-conv padding + // rewrite below needs them. + val weightShape = node.inputs.getOrNull(1)?.shape + if (weightShape == null || weightShape.size != 4) { + return ConversionResult.Failure( + "convTranspose2d lowering requires a static 4D weight shape, got $weightShape", + "Missing weight shape for node ${node.id}" + ) + } + + val outputSpec = node.outputs.firstOrNull() + val outputType = outputSpec?.let { context.getTypeMapper().mapTensorType(it) } + ?: "tensor" + val typeMapper = context.getTypeMapper() + val inputType = node.inputs.getOrNull(0)?.let { typeMapper.mapTensorType(it) } ?: "tensor" + val weightType = node.inputs.getOrNull(1)?.let { typeMapper.mapTensorType(it) } ?: "tensor" + val biasType = node.inputs.getOrNull(2)?.let { typeMapper.mapTensorType(it) } + + val resultValue = context.nextTempValue() + val operations = buildConvTranspose2dOperation( + resultValue = resultValue, + input = operands[0], + weight = operands[1], + bias = if (operands.size > 2) operands[2] else null, + inputType = inputType, + weightType = weightType, + outputType = outputType, + biasType = biasType, + kernelSize = weightShape[2] to weightShape[3], + stride = stride, + padding = padding, + outputPadding = outputPadding, + dilation = dilation, + ) + operations.forEach { context.emitOperation(it) } + + return ConversionResult.Success( + outputValueName = resultValue, + emittedOperations = operations + ) + } + private fun convertMaxPool2d( node: GraphNode, operands: List, @@ -1077,6 +1145,18 @@ public class NeuralNetOperationsConverter : StableHloOperationConverter { } } + private fun extractOutputPadding(params: Map): Pair { + return when (val outputPadding = params["outputPadding"] ?: params["output_padding"]) { + is Pair<*, *> -> (outputPadding.first as? Int ?: 0) to (outputPadding.second as? Int ?: 0) + is Int -> outputPadding to outputPadding + is List<*> -> { + val list = outputPadding.filterIsInstance() + if (list.size >= 2) list[0] to list[1] else 0 to 0 + } + else -> 0 to 0 + } + } + private fun extractKernelSize(params: Map): Pair { return when (val kernelSize = params["kernelSize"] ?: params["kernel_size"]) { is Pair<*, *> -> (kernelSize.first as? Int ?: 2) to (kernelSize.second as? Int ?: 2) @@ -1174,6 +1254,65 @@ public class NeuralNetOperationsConverter : StableHloOperationConverter { } } + /** + * Lower `convTranspose2d` to a regular `stablehlo.convolution` (the XLA conv-transpose + * rewrite): the kernel is spatially flipped (`stablehlo.reverse` on dims [2, 3]) and read + * with `[i, o, 0, 1]` dim numbers (weight layout is `[C_in, C_out, kH, kW]`), the stride + * becomes `lhs_dilate`, and the padding per axis becomes + * `low = dilation*(k-1) - padding`, `high = dilation*(k-1) - padding + outputPadding`. + */ + private fun buildConvTranspose2dOperation( + resultValue: String, + input: String, + weight: String, + bias: String?, + inputType: String, + weightType: String, + outputType: String, + biasType: String?, + kernelSize: Pair, + stride: Pair, + padding: Pair, + outputPadding: Pair, + dilation: Pair, + ): List { + val (kH, kW) = kernelSize + val padLowH = dilation.first * (kH - 1) - padding.first + val padHighH = padLowH + outputPadding.first + val padLowW = dilation.second * (kW - 1) - padding.second + val padHighW = padLowW + outputPadding.second + + val reversedWeight = "${resultValue}_wrev" + val reverseOp = "$reversedWeight = stablehlo.reverse $weight, dims = [2, 3] : $weightType" + + val convCore = { rv: String -> + "$rv = stablehlo.convolution($input, $reversedWeight) " + + "dim_numbers = [b, f, 0, 1]x[i, o, 0, 1]->[b, f, 0, 1], " + + "window = {stride = [1, 1], " + + "pad = [[$padLowH, $padHighH], [$padLowW, $padHighW]], " + + "lhs_dilate = [${stride.first}, ${stride.second}], " + + "rhs_dilate = [${dilation.first}, ${dilation.second}]} " + + "{batch_group_count = 1 : i64, feature_group_count = 1 : i64} " + + ": ($inputType, $weightType) -> $outputType" + } + + return if (bias != null) { + val convResult = "${resultValue}_conv" + val biasBcast = "${resultValue}_bias_b" + // bias is [Cout]; conv output is [N, Cout, H, W]; broadcast along feature dim (index 1) + val bcastOp = "$biasBcast = stablehlo.broadcast_in_dim $bias, dims = [1] : " + + "(${biasType ?: "tensor"}) -> $outputType" + listOf( + reverseOp, + convCore(convResult), + bcastOp, + "$resultValue = stablehlo.add $convResult, $biasBcast : $outputType" + ) + } else { + listOf(reverseOp, convCore(resultValue)) + } + } + /** The MLIR element type ("f32"/"f16"/โ€ฆ) parsed from a `tensor<โ€ฆxT>` string. */ private fun elementTypeOf(tensorType: String): String = tensorType.substringAfterLast('x').substringBefore('>').ifBlank { "f32" } diff --git a/skainet-compile/skainet-compile-hlo/src/commonTest/kotlin/sk/ainet/compile/hlo/ConvTranspose2dConverterTest.kt b/skainet-compile/skainet-compile-hlo/src/commonTest/kotlin/sk/ainet/compile/hlo/ConvTranspose2dConverterTest.kt new file mode 100644 index 000000000..a25fc152e --- /dev/null +++ b/skainet-compile/skainet-compile-hlo/src/commonTest/kotlin/sk/ainet/compile/hlo/ConvTranspose2dConverterTest.kt @@ -0,0 +1,140 @@ +package sk.ainet.compile.hlo + +import sk.ainet.compile.hlo.converters.NeuralNetOperationsConverter +import sk.ainet.lang.graph.DefaultComputeGraph +import sk.ainet.lang.graph.GraphEdge +import sk.ainet.lang.graph.GraphNode +import sk.ainet.lang.tensor.Tensor +import sk.ainet.lang.tensor.ops.Operation +import sk.ainet.lang.tensor.ops.TensorSpec +import sk.ainet.lang.tensor.ops.ValidationResult +import sk.ainet.lang.types.DType +import kotlin.test.Test +import kotlin.test.assertContains +import kotlin.test.assertFalse +import kotlin.test.assertTrue + +/** + * Covers the convTranspose2d StableHLO lowering (#1371 / #1375): a transposed + * convolution lowers to a regular `stablehlo.convolution` with the stride moved + * to `lhs_dilate`, the kernel spatially flipped (`stablehlo.reverse`) and read + * with `[i, o, 0, 1]` dim numbers, and the padding rewritten per axis to + * `low = dilation*(k-1) - padding`, `high = low + outputPadding`. + */ +class ConvTranspose2dConverterTest { + + @Test + fun convTranspose2d_operation_is_supported() { + val registry = StableHloOperationRegistry() + registry.register(NeuralNetOperationsConverter()) + assertTrue(registry.isSupported("convTranspose2d")) + } + + @Test + fun convTranspose2d_lowers_to_lhs_dilated_convolution_with_flipped_kernel() { + // input [1,2,3,3], weight [C_in=2, C_out=3, 3, 3], stride 2, padding 1, outputPadding 1 + // -> output [1,3,6,6]; pad low = 1*(3-1) - 1 = 1, pad high = 1 + 1 = 2. + val graph = buildConvTranspose2dGraph( + stride = 2 to 2, padding = 1 to 1, outputPadding = 1 to 1, dilation = 1 to 1, groups = 1, + ) + val converter = StableHloConverterFactory.createExtended() + val module = converter.convert(graph, "test_conv_transpose2d") + + assertFalse( + module.content.contains("Unsupported operation") || module.content.contains("No converter found"), + "convTranspose2d must be claimed by a converter:\n${module.content}" + ) + assertContains(module.content, "stablehlo.reverse", false, "kernel must be spatially flipped") + assertContains(module.content, "dims = [2, 3]", false, "flip must be over the spatial dims") + assertContains(module.content, "stablehlo.convolution") + assertContains(module.content, "[i, o, 0, 1]", false, "kernel dim numbers must swap in/out features") + assertContains(module.content, "lhs_dilate = [2, 2]", false, "stride must become lhs dilation") + assertContains(module.content, "stride = [1, 1]", false, "window stride must be 1") + assertContains(module.content, "pad = [[1, 2], [1, 2]]", false, "padding must be rewritten with outputPadding on the high side") + } + + @Test + fun convTranspose2d_grouped_lowering_is_rejected() { + val graph = buildConvTranspose2dGraph( + stride = 1 to 1, padding = 0 to 0, outputPadding = 0 to 0, dilation = 1 to 1, groups = 2, + ) + val converter = StableHloConverterFactory.createExtended() + val result = runCatching { converter.convert(graph, "test_conv_transpose2d_grouped") } + + // Either the conversion throws or the module records the failure โ€” it must not + // silently emit a groups-ignoring convolution. + val content = result.getOrNull()?.content + assertTrue( + result.isFailure || content == null || !content.contains("lhs_dilate"), + "grouped convTranspose2d must not silently lower as groups=1:\n$content" + ) + } + + private fun buildConvTranspose2dGraph( + stride: Pair, + padding: Pair, + outputPadding: Pair, + dilation: Pair, + groups: Int, + ): DefaultComputeGraph { + val graph = DefaultComputeGraph() + val inputSpec = TensorSpec("input", listOf(1, 2, 3, 3), "FP32") + val weightSpec = TensorSpec("weight", listOf(2, 3, 3, 3), "FP32") + val outH = (3 - 1) * stride.first - 2 * padding.first + dilation.first * 2 + outputPadding.first + 1 + val outW = (3 - 1) * stride.second - 2 * padding.second + dilation.second * 2 + outputPadding.second + 1 + val outputSpec = TensorSpec("out", listOf(1, 3 * groups, outH, outW), "FP32") + + val input = GraphNode("input", markerInputOp(), emptyList(), listOf(inputSpec)) + val weight = GraphNode("weight", markerInputOp(), emptyList(), listOf(weightSpec)) + val convT = GraphNode( + id = "convTranspose2d", + operation = convTranspose2dOp(stride, padding, outputPadding, dilation, groups), + inputs = listOf(inputSpec, weightSpec), + outputs = listOf(outputSpec) + ) + graph.addNode(input) + graph.addNode(weight) + graph.addNode(convT) + graph.addEdge(GraphEdge("e1", input, convT, 0, 0, inputSpec)) + graph.addEdge(GraphEdge("e2", weight, convT, 0, 1, weightSpec)) + return graph + } + + private fun markerInputOp(): Operation = object : Operation { + override val name: String = "input" + override val type: String = "input" + override val parameters: Map = emptyMap() + override fun execute(inputs: List>): List> = + throw UnsupportedOperationException("test fixture only") + override fun validateInputs(inputs: List): ValidationResult = ValidationResult.Valid + override fun inferOutputs(inputs: List): List = emptyList() + override fun clone(newParameters: Map): Operation = this + override fun serialize(): Map = mapOf("name" to name, "type" to type) + } + + private fun convTranspose2dOp( + stride: Pair, + padding: Pair, + outputPadding: Pair, + dilation: Pair, + groups: Int, + ): Operation = object : Operation { + override val name: String = "convTranspose2d" + override val type: String = "nn" + override val parameters: Map = mapOf( + "stride" to stride, + "padding" to padding, + "outputPadding" to outputPadding, + "dilation" to dilation, + "groups" to groups, + ) + override fun execute(inputs: List>): List> = + throw UnsupportedOperationException("test fixture only") + override fun validateInputs(inputs: List): ValidationResult = ValidationResult.Valid + override fun inferOutputs(inputs: List): List = inputs.take(1) + override fun clone(newParameters: Map): Operation = this + override fun serialize(): Map = mapOf( + "name" to name, "type" to type, "parameters" to parameters + ) + } +} diff --git a/skainet-lang/skainet-lang-core/api/jvm/skainet-lang-core.api b/skainet-lang/skainet-lang-core/api/jvm/skainet-lang-core.api index 7d4ed512a..89f1aabab 100644 --- a/skainet-lang/skainet-lang-core/api/jvm/skainet-lang-core.api +++ b/skainet-lang/skainet-lang-core/api/jvm/skainet-lang-core.api @@ -6408,6 +6408,7 @@ public abstract interface class sk/ainet/lang/tensor/ops/DifferentiableTensorOps public abstract fun conv2dBackward (Lsk/ainet/lang/tensor/Tensor;Lsk/ainet/lang/tensor/Tensor;Ljava/util/List;Ljava/util/Map;)Ljava/util/List; public abstract fun conv3dBackward (Lsk/ainet/lang/tensor/Tensor;Lsk/ainet/lang/tensor/Tensor;Ljava/util/List;Ljava/util/Map;)Ljava/util/List; public abstract fun convTranspose1dBackward (Lsk/ainet/lang/tensor/Tensor;Lsk/ainet/lang/tensor/Tensor;Ljava/util/List;Ljava/util/Map;)Ljava/util/List; + public abstract fun convTranspose2dBackward (Lsk/ainet/lang/tensor/Tensor;Lsk/ainet/lang/tensor/Tensor;Ljava/util/List;Ljava/util/Map;)Ljava/util/List; public abstract fun cosBackward (Lsk/ainet/lang/tensor/Tensor;Lsk/ainet/lang/tensor/Tensor;Ljava/util/List;Ljava/util/Map;)Ljava/util/List; public abstract fun divScalarBackward (Lsk/ainet/lang/tensor/Tensor;Lsk/ainet/lang/tensor/Tensor;Ljava/util/List;Ljava/util/Map;)Ljava/util/List; public abstract fun divideBackward (Lsk/ainet/lang/tensor/Tensor;Lsk/ainet/lang/tensor/Tensor;Ljava/util/List;Ljava/util/Map;)Ljava/util/List; @@ -6663,7 +6664,6 @@ public abstract interface class sk/ainet/lang/tensor/ops/MixedPrecisionTensorOps } public final class sk/ainet/lang/tensor/ops/MixedPrecisionTensorOps$DefaultImpls { - public static fun convTranspose2d (Lsk/ainet/lang/tensor/ops/MixedPrecisionTensorOps;Lsk/ainet/lang/tensor/Tensor;Lsk/ainet/lang/tensor/Tensor;Lsk/ainet/lang/tensor/Tensor;Lkotlin/Pair;Lkotlin/Pair;Lkotlin/Pair;Lkotlin/Pair;I)Lsk/ainet/lang/tensor/Tensor; public static fun matmulWeightTransposed (Lsk/ainet/lang/tensor/ops/MixedPrecisionTensorOps;Lsk/ainet/lang/tensor/Tensor;Lsk/ainet/lang/tensor/Tensor;)Lsk/ainet/lang/tensor/Tensor; public static fun relayoutPackedWeightForKernels (Lsk/ainet/lang/tensor/ops/MixedPrecisionTensorOps;Lsk/ainet/lang/tensor/Tensor;)Lsk/ainet/lang/tensor/Tensor; } @@ -6837,7 +6837,7 @@ public abstract interface class sk/ainet/lang/tensor/ops/TensorOps { public static synthetic fun conv3d$default (Lsk/ainet/lang/tensor/ops/TensorOps;Lsk/ainet/lang/tensor/Tensor;Lsk/ainet/lang/tensor/Tensor;Lsk/ainet/lang/tensor/Tensor;Lkotlin/Triple;Lkotlin/Triple;Lkotlin/Triple;IILjava/lang/Object;)Lsk/ainet/lang/tensor/Tensor; public abstract fun convTranspose1d (Lsk/ainet/lang/tensor/Tensor;Lsk/ainet/lang/tensor/Tensor;Lsk/ainet/lang/tensor/Tensor;IIIII)Lsk/ainet/lang/tensor/Tensor; public static synthetic fun convTranspose1d$default (Lsk/ainet/lang/tensor/ops/TensorOps;Lsk/ainet/lang/tensor/Tensor;Lsk/ainet/lang/tensor/Tensor;Lsk/ainet/lang/tensor/Tensor;IIIIIILjava/lang/Object;)Lsk/ainet/lang/tensor/Tensor; - public fun convTranspose2d (Lsk/ainet/lang/tensor/Tensor;Lsk/ainet/lang/tensor/Tensor;Lsk/ainet/lang/tensor/Tensor;Lkotlin/Pair;Lkotlin/Pair;Lkotlin/Pair;Lkotlin/Pair;I)Lsk/ainet/lang/tensor/Tensor; + public abstract fun convTranspose2d (Lsk/ainet/lang/tensor/Tensor;Lsk/ainet/lang/tensor/Tensor;Lsk/ainet/lang/tensor/Tensor;Lkotlin/Pair;Lkotlin/Pair;Lkotlin/Pair;Lkotlin/Pair;I)Lsk/ainet/lang/tensor/Tensor; public static synthetic fun convTranspose2d$default (Lsk/ainet/lang/tensor/ops/TensorOps;Lsk/ainet/lang/tensor/Tensor;Lsk/ainet/lang/tensor/Tensor;Lsk/ainet/lang/tensor/Tensor;Lkotlin/Pair;Lkotlin/Pair;Lkotlin/Pair;Lkotlin/Pair;IILjava/lang/Object;)Lsk/ainet/lang/tensor/Tensor; public abstract fun convert (Lsk/ainet/lang/tensor/Tensor;Lsk/ainet/lang/types/DType;)Lsk/ainet/lang/tensor/Tensor; public abstract fun cos (Lsk/ainet/lang/tensor/Tensor;)Lsk/ainet/lang/tensor/Tensor; @@ -6917,7 +6917,6 @@ public final class sk/ainet/lang/tensor/ops/TensorOps$DefaultImpls { public static synthetic fun conv2d$default (Lsk/ainet/lang/tensor/ops/TensorOps;Lsk/ainet/lang/tensor/Tensor;Lsk/ainet/lang/tensor/Tensor;Lsk/ainet/lang/tensor/Tensor;Lkotlin/Pair;Lkotlin/Pair;Lkotlin/Pair;IILjava/lang/Object;)Lsk/ainet/lang/tensor/Tensor; public static synthetic fun conv3d$default (Lsk/ainet/lang/tensor/ops/TensorOps;Lsk/ainet/lang/tensor/Tensor;Lsk/ainet/lang/tensor/Tensor;Lsk/ainet/lang/tensor/Tensor;Lkotlin/Triple;Lkotlin/Triple;Lkotlin/Triple;IILjava/lang/Object;)Lsk/ainet/lang/tensor/Tensor; public static synthetic fun convTranspose1d$default (Lsk/ainet/lang/tensor/ops/TensorOps;Lsk/ainet/lang/tensor/Tensor;Lsk/ainet/lang/tensor/Tensor;Lsk/ainet/lang/tensor/Tensor;IIIIIILjava/lang/Object;)Lsk/ainet/lang/tensor/Tensor; - public static fun convTranspose2d (Lsk/ainet/lang/tensor/ops/TensorOps;Lsk/ainet/lang/tensor/Tensor;Lsk/ainet/lang/tensor/Tensor;Lsk/ainet/lang/tensor/Tensor;Lkotlin/Pair;Lkotlin/Pair;Lkotlin/Pair;Lkotlin/Pair;I)Lsk/ainet/lang/tensor/Tensor; public static synthetic fun convTranspose2d$default (Lsk/ainet/lang/tensor/ops/TensorOps;Lsk/ainet/lang/tensor/Tensor;Lsk/ainet/lang/tensor/Tensor;Lsk/ainet/lang/tensor/Tensor;Lkotlin/Pair;Lkotlin/Pair;Lkotlin/Pair;Lkotlin/Pair;IILjava/lang/Object;)Lsk/ainet/lang/tensor/Tensor; public static synthetic fun elu$default (Lsk/ainet/lang/tensor/ops/TensorOps;Lsk/ainet/lang/tensor/Tensor;FILjava/lang/Object;)Lsk/ainet/lang/tensor/Tensor; public static synthetic fun flatten$default (Lsk/ainet/lang/tensor/ops/TensorOps;Lsk/ainet/lang/tensor/Tensor;IIILjava/lang/Object;)Lsk/ainet/lang/tensor/Tensor; diff --git a/skainet-lang/skainet-lang-core/src/commonMain/kotlin/sk/ainet/lang/nn/TransposedConv2d.kt b/skainet-lang/skainet-lang-core/src/commonMain/kotlin/sk/ainet/lang/nn/TransposedConv2d.kt index ddd8fbea7..896f6fe66 100644 --- a/skainet-lang/skainet-lang-core/src/commonMain/kotlin/sk/ainet/lang/nn/TransposedConv2d.kt +++ b/skainet-lang/skainet-lang-core/src/commonMain/kotlin/sk/ainet/lang/nn/TransposedConv2d.kt @@ -5,7 +5,6 @@ import sk.ainet.lang.tensor.Tensor import sk.ainet.lang.types.DType import sk.ainet.lang.nn.topology.ModuleParameter import sk.ainet.lang.nn.topology.ModuleParameters -import sk.ainet.lang.ops.NotImplemented /** * Transposed Convolution (Deconvolution) layer. @@ -37,7 +36,6 @@ import sk.ainet.lang.ops.NotImplemented * @param initWeights Initial weights tensor * @param initBias Initial bias tensor (if bias is true) */ -@NotImplemented(reason = "forward needs TensorOps.convTranspose2d, which no backend implements yet", issue = "#1371") public class TransposedConv2d( public val inChannels: Int, public val outChannels: Int, @@ -82,7 +80,6 @@ public class TransposedConv2d( sk.ainet.lang.nn.hooks.withForwardHooks(ctx, this, input) { val weight = params.first { it.name.endsWith(".weight") }.value val biasValue = if (bias) params.firstOrNull { it.name.endsWith(".bias") }?.value else null - // Throws until a backend implements convTranspose2d (see @NotImplemented above). input.ops.convTranspose2d(input, weight, biasValue, stride, padding, outputPadding, dilation, groups) } diff --git a/skainet-lang/skainet-lang-core/src/commonMain/kotlin/sk/ainet/lang/tensor/ops/TensorOps.kt b/skainet-lang/skainet-lang-core/src/commonMain/kotlin/sk/ainet/lang/tensor/ops/TensorOps.kt index 2ade46ad8..ad78125b8 100644 --- a/skainet-lang/skainet-lang-core/src/commonMain/kotlin/sk/ainet/lang/tensor/ops/TensorOps.kt +++ b/skainet-lang/skainet-lang-core/src/commonMain/kotlin/sk/ainet/lang/tensor/ops/TensorOps.kt @@ -9,7 +9,6 @@ import sk.ainet.lang.nn.dsl.GenerateNetworkDsl import sk.ainet.lang.nn.dsl.ActivationDsl import sk.ainet.lang.ops.DarcValidated import sk.ainet.lang.ops.Lowering -import sk.ainet.lang.ops.NotImplemented @GenerateTracingWrapper @GenerateNetworkDsl @@ -165,11 +164,8 @@ public interface TensorOps { * * [weight] is laid out `[C_in, C_out / groups, kH, kW]` (PyTorch `ConvTranspose2d`). Output * spatial size per axis is `(in - 1) * stride - 2 * padding + dilation * (k - 1) + outputPadding + 1`. - * - * Not implemented on any backend yet: the default throws. Once a backend kernel lands this - * becomes an abstract `@Diff` op like [convTranspose1d]. */ - @NotImplemented(reason = "no backend kernel yet; TransposedConv2d.forward depends on it", issue = "#1371") + @Diff public fun convTranspose2d( input: Tensor, weight: Tensor, @@ -179,7 +175,7 @@ public interface TensorOps { outputPadding: Pair = 0 to 0, dilation: Pair = 1 to 1, groups: Int = 1 - ): Tensor = throw NotImplementedError("convTranspose2d is not implemented on this backend yet (see #1371)") + ): Tensor // Pooling operations @Diff diff --git a/skainet-lang/skainet-lang-core/src/commonMain/kotlin/sk/ainet/lang/tensor/ops/VoidTensorOps.kt b/skainet-lang/skainet-lang-core/src/commonMain/kotlin/sk/ainet/lang/tensor/ops/VoidTensorOps.kt index 8ca0194a6..da22db57e 100644 --- a/skainet-lang/skainet-lang-core/src/commonMain/kotlin/sk/ainet/lang/tensor/ops/VoidTensorOps.kt +++ b/skainet-lang/skainet-lang-core/src/commonMain/kotlin/sk/ainet/lang/tensor/ops/VoidTensorOps.kt @@ -246,6 +246,21 @@ public class VoidTensorOps : TensorOps { return VoidOpsTensor(resultData, input.dtype) } + override fun convTranspose2d( + input: Tensor, + weight: Tensor, + bias: Tensor?, + stride: Pair, + padding: Pair, + outputPadding: Pair, + dilation: Pair, + groups: Int + ): Tensor { + val resultShape = calculateConvTranspose2dShape(input.shape, weight.shape, stride, padding, outputPadding, dilation, groups) + val resultData = dataFactory.zeros(resultShape, input.dtype) + return VoidOpsTensor(resultData, input.dtype) + } + override fun maxPool2d( input: Tensor, kernelSize: Pair, @@ -913,6 +928,30 @@ public class VoidTensorOps : TensorOps { return Shape(batch, outChannels, outputLength) } + /** + * Calculates the result shape for convTranspose2d operation. + * Input shape: (batch, in_channels, height, width) + * Weight shape: (in_channels, out_channels_per_group, kernel_height, kernel_width) + * Output shape: (batch, out_channels, out_height, out_width) + */ + private fun calculateConvTranspose2dShape( + inputShape: Shape, + weightShape: Shape, + stride: Pair, + padding: Pair, + outputPadding: Pair, + dilation: Pair, + groups: Int + ): Shape { + val batch = inputShape[0] + val outChannels = weightShape[1] * groups + val outputHeight = + (inputShape[2] - 1) * stride.first - 2 * padding.first + dilation.first * (weightShape[2] - 1) + outputPadding.first + 1 + val outputWidth = + (inputShape[3] - 1) * stride.second - 2 * padding.second + dilation.second * (weightShape[3] - 1) + outputPadding.second + 1 + return Shape(batch, outChannels, outputHeight, outputWidth) + } + /** * Calculates the result shape for conv2d operation. * Input shape: (batch, in_channels, height, width) diff --git a/skainet-lang/skainet-lang-dag/api/jvm/skainet-lang-dag.api b/skainet-lang/skainet-lang-dag/api/jvm/skainet-lang-dag.api index 4969d8fec..c031a45f1 100644 --- a/skainet-lang/skainet-lang-dag/api/jvm/skainet-lang-dag.api +++ b/skainet-lang/skainet-lang-dag/api/jvm/skainet-lang-dag.api @@ -36,7 +36,6 @@ public abstract interface class sk/ainet/lang/dag/GraphDslOps : sk/ainet/lang/te } public final class sk/ainet/lang/dag/GraphDslOps$DefaultImpls { - public static fun convTranspose2d (Lsk/ainet/lang/dag/GraphDslOps;Lsk/ainet/lang/tensor/Tensor;Lsk/ainet/lang/tensor/Tensor;Lsk/ainet/lang/tensor/Tensor;Lkotlin/Pair;Lkotlin/Pair;Lkotlin/Pair;Lkotlin/Pair;I)Lsk/ainet/lang/tensor/Tensor; public static fun matmulWeightTransposed (Lsk/ainet/lang/dag/GraphDslOps;Lsk/ainet/lang/tensor/Tensor;Lsk/ainet/lang/tensor/Tensor;)Lsk/ainet/lang/tensor/Tensor; public static fun relayoutPackedWeightForKernels (Lsk/ainet/lang/dag/GraphDslOps;Lsk/ainet/lang/tensor/Tensor;)Lsk/ainet/lang/tensor/Tensor; } @@ -64,6 +63,8 @@ public final class sk/ainet/lang/dag/GraphDslOpsGraphDslKt { public static synthetic fun conv3d$default (Lsk/ainet/lang/dag/DagBuilder;Lsk/ainet/lang/dag/GraphValue;Lsk/ainet/lang/dag/GraphValue;Lsk/ainet/lang/dag/GraphValue;Lkotlin/Triple;Lkotlin/Triple;Lkotlin/Triple;ILjava/lang/String;ILjava/lang/Object;)Lsk/ainet/lang/dag/GraphValue; public static final fun convTranspose1d (Lsk/ainet/lang/dag/DagBuilder;Lsk/ainet/lang/dag/GraphValue;Lsk/ainet/lang/dag/GraphValue;Lsk/ainet/lang/dag/GraphValue;IIIIILjava/lang/String;)Lsk/ainet/lang/dag/GraphValue; public static synthetic fun convTranspose1d$default (Lsk/ainet/lang/dag/DagBuilder;Lsk/ainet/lang/dag/GraphValue;Lsk/ainet/lang/dag/GraphValue;Lsk/ainet/lang/dag/GraphValue;IIIIILjava/lang/String;ILjava/lang/Object;)Lsk/ainet/lang/dag/GraphValue; + public static final fun convTranspose2d (Lsk/ainet/lang/dag/DagBuilder;Lsk/ainet/lang/dag/GraphValue;Lsk/ainet/lang/dag/GraphValue;Lsk/ainet/lang/dag/GraphValue;Lkotlin/Pair;Lkotlin/Pair;Lkotlin/Pair;Lkotlin/Pair;ILjava/lang/String;)Lsk/ainet/lang/dag/GraphValue; + public static synthetic fun convTranspose2d$default (Lsk/ainet/lang/dag/DagBuilder;Lsk/ainet/lang/dag/GraphValue;Lsk/ainet/lang/dag/GraphValue;Lsk/ainet/lang/dag/GraphValue;Lkotlin/Pair;Lkotlin/Pair;Lkotlin/Pair;Lkotlin/Pair;ILjava/lang/String;ILjava/lang/Object;)Lsk/ainet/lang/dag/GraphValue; public static final fun cos (Lsk/ainet/lang/dag/DagBuilder;Lsk/ainet/lang/dag/GraphValue;Ljava/lang/String;)Lsk/ainet/lang/dag/GraphValue; public static synthetic fun cos$default (Lsk/ainet/lang/dag/DagBuilder;Lsk/ainet/lang/dag/GraphValue;Ljava/lang/String;ILjava/lang/Object;)Lsk/ainet/lang/dag/GraphValue; public static final fun divScalar (Lsk/ainet/lang/dag/DagBuilder;Lsk/ainet/lang/dag/GraphValue;Ljava/lang/Number;Ljava/lang/String;)Lsk/ainet/lang/dag/GraphValue; diff --git a/skainet-lang/skainet-lang-ksp-processor/src/main/kotlin/sk/ainet/lang/ksp/TracingWrapperProcessor.kt b/skainet-lang/skainet-lang-ksp-processor/src/main/kotlin/sk/ainet/lang/ksp/TracingWrapperProcessor.kt index 1784d2448..3a01caf11 100644 --- a/skainet-lang/skainet-lang-ksp-processor/src/main/kotlin/sk/ainet/lang/ksp/TracingWrapperProcessor.kt +++ b/skainet-lang/skainet-lang-ksp-processor/src/main/kotlin/sk/ainet/lang/ksp/TracingWrapperProcessor.kt @@ -335,7 +335,7 @@ class TracingWrapperGenerator( // Fix the type for nullable parameters that should have ? suffix val correctedType = when { param.name == "dim" && method.name in listOf("squeeze", "sum", "mean", "variance") && !param.type.endsWith("?") -> "${param.type}?" - param.name == "bias" && method.name in listOf("conv1d", "conv2d", "conv3d", "convTranspose1d") && !param.type.endsWith("?") -> "${param.type}?" + param.name == "bias" && method.name in listOf("conv1d", "conv2d", "conv3d", "convTranspose1d", "convTranspose2d") && !param.type.endsWith("?") -> "${param.type}?" param.name in listOf("mask", "keyLengths") && method.name == "scaledDotProductAttention" && !param.type.endsWith("?") -> "${param.type}?" else -> param.type }