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 42e50c73..96d52d15 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 @@ -484,6 +484,74 @@ public open class DefaultCpuOpsBase(protected val dataFactory: TensorDataFactory return newTensor(outData, tensor.dtype, tensor) } + @TensorOp() + override fun permute(tensor: Tensor, axes: IntArray): Tensor { + val rank = tensor.shape.rank + require(axes.size == rank) { + "permute: axes length ${axes.size} must match tensor rank $rank" + } + val seen = BooleanArray(rank) + for (a in axes) { + require(a in 0 until rank) { "permute: axis $a out of range [0, $rank)" } + require(!seen[a]) { "permute: axis $a appears more than once in ${axes.toList()}" } + seen[a] = true + } + + val inDims = tensor.shape.dimensions + val outDims = IntArray(rank) { i -> inDims[axes[i]] } + val outShape = Shape(outDims) + + // Identity permute — no copy. + var isIdentity = true + for (i in 0 until rank) if (axes[i] != i) { isIdentity = false; break } + if (isIdentity) return tensor + + // Row-major strides for input and output. inStrides[k] is the + // distance in the source buffer between consecutive indices on + // input axis k. + val inStrides = IntArray(rank).also { s -> + s[rank - 1] = 1 + for (i in rank - 2 downTo 0) s[i] = s[i + 1] * inDims[i + 1] + } + val outStrides = IntArray(rank).also { s -> + s[rank - 1] = 1 + for (i in rank - 2 downTo 0) s[i] = s[i + 1] * outDims[i + 1] + } + + // Fast path: source is a contiguous FloatArray. Iterate the output + // linearly, decompose each flat index to its multi-index, permute + // to source coords, recompose to source flat index, copy. + if (tensor.data is FloatArrayTensorData<*>) { + val srcBuf = (tensor.data as FloatArrayTensorData<*>).buffer + val total = outShape.volume + val out = FloatArray(total) + val outIdx = IntArray(rank) + for (flatOut in 0 until total) { + var rem = flatOut + for (i in 0 until rank) { + val s = outStrides[i] + outIdx[i] = rem / s + rem -= outIdx[i] * s + } + var flatIn = 0 + for (i in 0 until rank) flatIn += outIdx[i] * inStrides[axes[i]] + out[flatOut] = srcBuf[flatIn] + } + @Suppress("UNCHECKED_CAST") + val outData = dataFactory.fromFloatArray(outShape, tensor.dtype, out) + as sk.ainet.lang.tensor.data.TensorData + return newTensor(outData, tensor.dtype, tensor) + } + + // Generic fallback: defer to dataFactory.init with element access. + val outData = dataFactory.init(outShape, tensor.dtype) { outIdx -> + val inIdx = IntArray(rank) + for (i in 0 until rank) inIdx[axes[i]] = outIdx[i] + tensor.data.get(*inIdx) + } + return newTensor(outData, tensor.dtype, tensor) + } + @TensorOp() @InProgress("cpu", owner = "team:cpu", issue = "task-ops.md#op-conv2d") override fun conv2d( diff --git a/skainet-backends/skainet-backend-cpu/src/commonTest/kotlin/sk/ainet/exec/tensor/ops/PermuteTest.kt b/skainet-backends/skainet-backend-cpu/src/commonTest/kotlin/sk/ainet/exec/tensor/ops/PermuteTest.kt new file mode 100644 index 00000000..c8492fcb --- /dev/null +++ b/skainet-backends/skainet-backend-cpu/src/commonTest/kotlin/sk/ainet/exec/tensor/ops/PermuteTest.kt @@ -0,0 +1,124 @@ +package sk.ainet.exec.tensor.ops + +import kotlin.test.Test +import kotlin.test.assertContentEquals +import kotlin.test.assertEquals +import kotlin.test.assertFailsWith +import kotlin.test.assertSame +import sk.ainet.context.DirectCpuExecutionContext +import sk.ainet.lang.tensor.Shape +import sk.ainet.lang.types.FP32 + +class PermuteTest { + + private fun ctx() = DirectCpuExecutionContext() + + @Test + fun identityPermuteReturnsSameTensor() { + val ctx = ctx() + val t = ctx.fromFloatArray( + Shape(2, 3, 4), FP32::class, + FloatArray(24) { it.toFloat() } + ) + val out = ctx.ops.permute(t, intArrayOf(0, 1, 2)) + assertSame(t, out, "identity permute should return the input tensor") + } + + @Test + fun swapDim0AndDim1OnRank3() { + val ctx = ctx() + // Shape [A=2, B=3, C=4], elements 0..23 row-major. + // Element (a, b, c) flat = a*12 + b*4 + c. + val src = FloatArray(24) { it.toFloat() } + val t = ctx.fromFloatArray(Shape(2, 3, 4), FP32::class, src) + val out = ctx.ops.permute(t, intArrayOf(1, 0, 2)) + assertContentEquals(intArrayOf(3, 2, 4), out.shape.dimensions, "expected shape [B=3, A=2, C=4]") + // out(b, a, c) == in(a, b, c) + for (b in 0 until 3) { + for (a in 0 until 2) { + for (c in 0 until 4) { + val expected = (a * 12 + b * 4 + c).toFloat() + val actual = out.data.get(b, a, c) + assertEquals(expected, actual, "out[$b,$a,$c] vs in[$a,$b,$c]") + } + } + } + } + + @Test + fun reverseAxesOnRank4() { + val ctx = ctx() + // Shape [2, 3, 4, 5]. Permute (3, 2, 1, 0) → reverses all axes. + val src = FloatArray(2 * 3 * 4 * 5) { it.toFloat() } + val t = ctx.fromFloatArray(Shape(2, 3, 4, 5), FP32::class, src) + val out = ctx.ops.permute(t, intArrayOf(3, 2, 1, 0)) + assertContentEquals(intArrayOf(5, 4, 3, 2), out.shape.dimensions) + for (d in 0 until 5) { + for (c in 0 until 4) { + for (b in 0 until 3) { + for (a in 0 until 2) { + val flatIn = a * 60 + b * 20 + c * 5 + d + assertEquals( + flatIn.toFloat(), + out.data.get(d, c, b, a), + "out[$d,$c,$b,$a] vs in[$a,$b,$c,$d]" + ) + } + } + } + } + } + + @Test + fun roundTripPermuteIsIdentity() { + val ctx = ctx() + val src = FloatArray(2 * 3 * 4) { it.toFloat() } + val t = ctx.fromFloatArray(Shape(2, 3, 4), FP32::class, src) + val axes = intArrayOf(2, 0, 1) + val inverse = IntArray(3).also { for (i in axes.indices) it[axes[i]] = i } + + val once = ctx.ops.permute(t, axes) + val back = ctx.ops.permute(once, inverse) + + assertContentEquals(t.shape.dimensions, back.shape.dimensions) + for (a in 0 until 2) for (b in 0 until 3) for (c in 0 until 4) { + assertEquals(t.data.get(a, b, c), back.data.get(a, b, c), "round-trip mismatch at [$a,$b,$c]") + } + } + + @Test + fun permuteEquivalentToTransposeOnRank2() { + val ctx = ctx() + val t = ctx.fromFloatArray( + Shape(3, 5), FP32::class, + FloatArray(15) { it.toFloat() } + ) + val viaPermute = ctx.ops.permute(t, intArrayOf(1, 0)) + val viaTranspose = ctx.ops.transpose(t) + assertContentEquals(viaTranspose.shape.dimensions, viaPermute.shape.dimensions) + for (i in 0 until 5) for (j in 0 until 3) { + assertEquals(viaTranspose.data.get(i, j), viaPermute.data.get(i, j)) + } + } + + @Test + fun rejectsWrongAxesLength() { + val ctx = ctx() + val t = ctx.fromFloatArray(Shape(2, 3), FP32::class, FloatArray(6)) + assertFailsWith { ctx.ops.permute(t, intArrayOf(1, 0, 2)) } + } + + @Test + fun rejectsOutOfRangeAxis() { + val ctx = ctx() + val t = ctx.fromFloatArray(Shape(2, 3), FP32::class, FloatArray(6)) + assertFailsWith { ctx.ops.permute(t, intArrayOf(0, 5)) } + } + + @Test + fun rejectsDuplicateAxis() { + val ctx = ctx() + val t = ctx.fromFloatArray(Shape(2, 3), FP32::class, FloatArray(6)) + assertFailsWith { ctx.ops.permute(t, intArrayOf(0, 0)) } + } +} 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 15320c62..c201d513 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 @@ -234,6 +234,13 @@ internal class RecordingTensorOpsDecorator(private val base: TensorOps) : Tensor return out } + override fun permute(tensor: Tensor, axes: IntArray): Tensor { + // Record as a regular passthrough; permute is shape-only at the + // op level. A dedicated PermuteOperation can be introduced later + // if the tape consumer needs to distinguish it from raw passthrough. + return base.permute(tensor, axes) + } + // --- Conv/Pool --- override fun conv1d( input: Tensor, 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 3b2bbfb2..1e70aba1 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 @@ -605,6 +605,16 @@ public class DefaultGradientTape( override fun transposeBackward(upstream: Tensor, output: Tensor, inputs: List>, attributes: Map): List?> = listOf(upstream.ops.transpose(upstream)) + override fun permuteBackward(upstream: Tensor, output: Tensor, inputs: List>, attributes: Map): List?> { + // Gradient of permute(t, axes) is permute(upstream, inverseAxes) + // where inverseAxes[axes[i]] = i. + val axes = (attributes["axes"] as? IntArray) + ?: error("permuteBackward: missing 'axes' attribute") + val inverse = IntArray(axes.size) + for (i in axes.indices) inverse[axes[i]] = i + return listOf(upstream.ops.permute(upstream, inverse)) + } + override fun reluBackward(upstream: Tensor, output: Tensor, inputs: List>, attributes: Map): List?> = listOf(reluGrad(upstream, inputs[0], output)) 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 8cace5ad..63406465 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 @@ -164,6 +164,7 @@ private class TestTensorOps : TensorOps { override fun rdivScalar(a: Number, b: Tensor): Tensor = b override fun matmul(a: Tensor, b: Tensor): Tensor = a override fun transpose(tensor: Tensor): Tensor = tensor + override fun permute(tensor: Tensor, axes: IntArray): Tensor = tensor override fun relu(tensor: Tensor): Tensor = tensor override fun leakyRelu(tensor: Tensor, negativeSlope: Float): Tensor = tensor override fun elu(tensor: Tensor, alpha: Float): Tensor = tensor 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 47c0001d..aad15c67 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 @@ -53,6 +53,24 @@ public interface TensorOps { @Diff public fun transpose(tensor: Tensor): Tensor + /** + * Permute the dimensions of [tensor] according to [axes]. + * + * `axes` is a permutation of `0..tensor.rank-1`; the i-th axis of the + * result is the `axes[i]`-th axis of the input. On a rank-3 tensor of + * shape `[A, B, C]`, `permute(t, intArrayOf(1, 0, 2))` returns shape + * `[B, A, C]`. + * + * `permute(t, intArrayOf(0, 1, ..., rank-3, rank-1, rank-2))` is + * equivalent to [transpose]. + * + * @param tensor input tensor, any rank ≥ 1 + * @param axes a permutation of `0..tensor.rank-1` (length must equal + * `tensor.rank`, every value in `[0, rank)` exactly once) + */ + @Diff + public fun permute(tensor: Tensor, axes: IntArray): Tensor + // Convolutional operations @Diff public fun conv1d( 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 8530d449..b4db5342 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 @@ -151,6 +151,13 @@ public class VoidTensorOps : TensorOps { return VoidOpsTensor(resultData, tensor.dtype) } + override fun permute(tensor: Tensor, axes: IntArray): Tensor { + validatePermuteAxes(tensor.shape, axes) + val resultShape = calculatePermuteShape(tensor.shape, axes) + val resultData = dataFactory.zeros(resultShape, tensor.dtype) + return VoidOpsTensor(resultData, tensor.dtype) + } + override fun conv1d( input: Tensor, weight: Tensor, @@ -598,6 +605,31 @@ public class VoidTensorOps : TensorOps { * For 2D tensors: (m, n) -> (n, m) * For higher dimensions: swaps the last two dimensions */ + /** + * Validate that [axes] is a valid permutation of `0..shape.rank-1`. + */ + internal fun validatePermuteAxes(shape: Shape, axes: IntArray) { + require(axes.size == shape.rank) { + "permute: axes length ${axes.size} must match tensor rank ${shape.rank}" + } + val seen = BooleanArray(shape.rank) + for (a in axes) { + require(a in 0 until shape.rank) { + "permute: axis $a out of range [0, ${shape.rank})" + } + require(!seen[a]) { "permute: axis $a appears more than once in $axes" } + seen[a] = true + } + } + + /** + * Result shape after applying [axes] permutation to [shape]. + */ + internal fun calculatePermuteShape(shape: Shape, axes: IntArray): Shape { + val dims = IntArray(shape.rank) { i -> shape.dimensions[axes[i]] } + return Shape(dims) + } + private fun calculateTransposeShape(shape: Shape): Shape { if (shape.rank < 2) { throw IllegalArgumentException("Transpose requires tensors with at least 2 dimensions")