Skip to content

feat(io-onnx): import ONNX ConvTranspose as real convTranspose1d/2d ops #1468

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
michalharakal merged 1 commit into from
Nov 5, 2026
Merged

feat(io-onnx): import ONNX ConvTranspose as real convTranspose1d/2d ops #1468

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
43 changes: 43 additions & 0 deletions ...kainet-compile-dag/src/commonMain/kotlin/sk/ainet/lang/graph/exec/ComputeGraphExecutor.kt
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -202,6 +202,32 @@ public class ComputeGraphExecutor(
listOf(ops.add(ops.multiply(normalized, scale), bias))
}

// Convolution ops
"convTranspose1d" -> listOf(
ops.convTranspose1d(
input = inputs[0],
weight = inputs[1],
bias = inputs.getOrNull(2),
stride = (params["stride"] as? Number)?.toInt() ?: 1,
padding = (params["padding"] as? Number)?.toInt() ?: 0,
outputPadding = (params["outputPadding"] as? Number)?.toInt() ?: 0,
dilation = (params["dilation"] as? Number)?.toInt() ?: 1,
groups = (params["groups"] as? Number)?.toInt() ?: 1,
)
)
"convTranspose2d" -> listOf(
ops.convTranspose2d(
input = inputs[0],
weight = inputs[1],
bias = inputs.getOrNull(2),
stride = params.intPair("stride", 1 to 1),
padding = params.intPair("padding", 0 to 0),
outputPadding = params.intPair("outputPadding", 0 to 0),
dilation = params.intPair("dilation", 1 to 1),
groups = (params["groups"] as? Number)?.toInt() ?: 1,
)
)

// Binary math ops
"add" -> listOf(ops.add(inputs[0], inputs[1]))
"subtract" -> listOf(ops.subtract(inputs[0], inputs[1]))
Expand Down Expand Up @@ -451,6 +477,23 @@ public class ComputeGraphExecutor(
}
}

/**
* Reads a per-axis Int pair recorded either as a `Pair` (direct op parameters), a two-element
* list (KSP traces and the ONNX importer), or a single number applied to both axes.
*/
private fun Map<String, Any>.intPair(key: String, default: Pair<Int, Int>): Pair<Int, Int> =
when (val value = this[key]) {
is Pair<*, *> -> ((value.first as? Number)?.toInt() ?: default.first) to
((value.second as? Number)?.toInt() ?: default.second)
is List<*> -> if (value.size == 2 && value.all { it is Number }) {
(value[0] as Number).toInt() to (value[1] as Number).toInt()
} else {
default
}
is Number -> value.toInt() to value.toInt()
else -> default
}

/**
* Handler for executing a fused operation.
*
Expand Down
134 changes: 134 additions & 0 deletions skainet-io/skainet-io-onnx/src/commonMain/kotlin/sk/ainet/io/onnx/OnnxImport.kt
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -244,6 +244,54 @@ public class OnnxToComputeGraphImporter {
attrInts("kernel_shape")?.let { put("kernelShape", it) }
}
)
"CONVTRANSPOSE" -> {
// ONNX weight layout equals TensorOps.convTranspose{1,2}d ([C_in, C_out/groups, k...]).
// ONNX pads are [*_begins..., *_ends...] per spatial axis while SKaiNET takes one
// symmetric padding per axis, so asymmetric pads (and auto_pad/output_shape, which
// derive pads at runtime) keep the placeholder and fail loudly on execute.
val strides = attrInts("strides")
val pads = attrInts("pads")
val dilations = attrInts("dilations")
val outputPadding = attrInts("output_padding")
val group = attributes["group"]?.ints?.firstOrNull()?.toSafeInt() ?: 1
val autoPad = attributes["auto_pad"]?.strings?.firstOrNull() ?: "NOTSET"
val spatialRank = attrInts("kernel_shape")?.size
?: strides?.size
?: dilations?.size
?: pads?.let { it.size / 2 }
?: inputs.getOrNull(1)?.takeIf { it.isNotEmpty() }
?.let { initializers[it]?.shape?.size?.minus(2) }
val symmetricPads = pads == null ||
(pads.size == 2 * (spatialRank ?: 0) &&
(0 until pads.size / 2).all { pads[it] == pads[it + pads.size / 2] })
val supported = autoPad == "NOTSET" && "output_shape" !in attributes &&
symmetricPads && (spatialRank == 1 || spatialRank == 2)
when {
!supported -> OnnxPlaceholderOperation(
opType = opType,
type = "onnx",
parameters = mapOf(
"reason" to "ConvTranspose supports 1D/2D with auto_pad=NOTSET, " +
"symmetric pads and no output_shape",
"autoPad" to autoPad,
)
)
spatialRank == 1 -> OnnxConvTranspose1dOperation(
stride = strides?.get(0) ?: 1,
padding = pads?.get(0) ?: 0,
outputPadding = outputPadding?.get(0) ?: 0,
dilation = dilations?.get(0) ?: 1,
groups = group,
)
else -> OnnxConvTranspose2dOperation(
stride = strides ?: listOf(1, 1),
padding = pads?.take(2) ?: listOf(0, 0),
outputPadding = outputPadding ?: listOf(0, 0),
dilation = dilations ?: listOf(1, 1),
groups = group,
)
}
}
"RELU" -> ReluOperation<DType, Any>()
"SIGMOID" -> SigmoidOperation<DType, Any>()
"LEAKYRELU" -> OnnxMappedOperation(
Expand Down Expand Up @@ -415,6 +463,8 @@ private fun opShapeInference(op: Operation, inputs: List<TensorSpec>, outputInde
is OnnxSliceOperation -> op.inferOutputs(inputs).getOrNull(outputIndex)?.shape
is OnnxSplitOperation -> op.inferOutputs(inputs).getOrNull(outputIndex)?.shape
is OnnxPadOperation -> op.inferOutputs(inputs).getOrNull(outputIndex)?.shape
is OnnxConvTranspose1dOperation -> op.inferOutputs(inputs).getOrNull(outputIndex)?.shape
is OnnxConvTranspose2dOperation -> op.inferOutputs(inputs).getOrNull(outputIndex)?.shape
is Conv2dOperation<*, *> -> inputs.firstOrNull()?.shape
is MaxPool2dOperation<*, *> -> inputs.firstOrNull()?.shape
is Upsample2dOperation<*, *> -> op.inferOutputs(inputs).getOrNull(outputIndex)?.shape
Expand Down Expand Up @@ -568,6 +618,90 @@ private class OnnxPadOperation(
}
}

/**
* ONNX `ConvTranspose` (1D) mapped to SKaiNET `convTranspose1d`. Output length is
* `(in - 1) * stride - 2 * padding + dilation * (k - 1) + outputPadding + 1`.
*/
private class OnnxConvTranspose1dOperation(
private val stride: Int,
private val padding: Int,
private val outputPadding: Int,
private val dilation: Int,
private val groups: Int,
) : BasePlaceholderOperation(
"convTranspose1d",
mapOf(
"stride" to stride,
"padding" to padding,
"outputPadding" to outputPadding,
"dilation" to dilation,
"groups" to groups,
),
"convolution",
) {
override fun inferOutputs(inputs: List<TensorSpec>): List<TensorSpec> {
val x = inputs.getOrNull(0)?.shape
val w = inputs.getOrNull(1)?.shape
val outputShape = if (x != null && w != null && x.size == 3 && w.size == 3) {
listOf(
x[0],
w[1] * groups,
(x[2] - 1) * stride - 2 * padding + dilation * (w[2] - 1) + outputPadding + 1,
)
} else {
x
}
return listOf(
TensorSpec(
name = "convTranspose1d_output",
shape = outputShape,
dtype = inputs.firstOrNull()?.dtype ?: "unknown"
)
)
}
}

/** ONNX `ConvTranspose` (2D) mapped to SKaiNET `convTranspose2d`; same output formula per H/W axis. */
private class OnnxConvTranspose2dOperation(
private val stride: List<Int>,
private val padding: List<Int>,
private val outputPadding: List<Int>,
private val dilation: List<Int>,
private val groups: Int,
) : BasePlaceholderOperation(
"convTranspose2d",
mapOf(
"stride" to stride,
"padding" to padding,
"outputPadding" to outputPadding,
"dilation" to dilation,
"groups" to groups,
),
"convolution",
) {
override fun inferOutputs(inputs: List<TensorSpec>): List<TensorSpec> {
val x = inputs.getOrNull(0)?.shape
val w = inputs.getOrNull(1)?.shape
val outputShape = if (x != null && w != null && x.size == 4 && w.size == 4) {
listOf(
x[0],
w[1] * groups,
(x[2] - 1) * stride[0] - 2 * padding[0] + dilation[0] * (w[2] - 1) + outputPadding[0] + 1,
(x[3] - 1) * stride[1] - 2 * padding[1] + dilation[1] * (w[3] - 1) + outputPadding[1] + 1,
)
} else {
x
}
return listOf(
TensorSpec(
name = "convTranspose2d_output",
shape = outputShape,
dtype = inputs.firstOrNull()?.dtype ?: "unknown"
)
)
}
}

private class OnnxPlaceholderOperation(
opType: String,
type: String = "onnx",
Expand Down
Loading
Loading