feat(ops): add TensorOps.permute(axes) for arbitrary-axis permutation

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 <T : DType, V> permute(tensor: Tensor<T, V>, axes: IntArray): Tensor<T, V> {
+        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<T, Float>(outShape, tensor.dtype, out)
+                as sk.ainet.lang.tensor.data.TensorData<T, V>
+            return newTensor(outData, tensor.dtype, tensor)
+        }
+
+        // Generic fallback: defer to dataFactory.init with element access.
+        val outData = dataFactory.init<T, V>(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 <T : DType, V> conv2d(

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<FP32, Float>(
+            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<FP32, Float>(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<FP32, Float>(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<FP32, Float>(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<FP32, Float>(
+            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<FP32, Float>(Shape(2, 3), FP32::class, FloatArray(6))
+        assertFailsWith<IllegalArgumentException> { ctx.ops.permute(t, intArrayOf(1, 0, 2)) }
+    }
+
+    @Test
+    fun rejectsOutOfRangeAxis() {
+        val ctx = ctx()
+        val t = ctx.fromFloatArray<FP32, Float>(Shape(2, 3), FP32::class, FloatArray(6))
+        assertFailsWith<IllegalArgumentException> { ctx.ops.permute(t, intArrayOf(0, 5)) }
+    }
+
+    @Test
+    fun rejectsDuplicateAxis() {
+        val ctx = ctx()
+        val t = ctx.fromFloatArray<FP32, Float>(Shape(2, 3), FP32::class, FloatArray(6))
+        assertFailsWith<IllegalArgumentException> { ctx.ops.permute(t, intArrayOf(0, 0)) }
+    }
+}

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 <T : DType, V> permute(tensor: Tensor<T, V>, axes: IntArray): Tensor<T, V> {
+        // 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 <T : DType, V> conv1d(
         input: Tensor<T, V>,

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<DType, Any>, output: Tensor<DType, Any>, inputs: List<Tensor<DType, Any>>, attributes: Map<String, Any?>): List<Tensor<DType, Any>?> =
         listOf(upstream.ops.transpose(upstream))
 
+    override fun permuteBackward(upstream: Tensor<DType, Any>, output: Tensor<DType, Any>, inputs: List<Tensor<DType, Any>>, attributes: Map<String, Any?>): List<Tensor<DType, Any>?> {
+        // 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<DType, Any>, output: Tensor<DType, Any>, inputs: List<Tensor<DType, Any>>, attributes: Map<String, Any?>): List<Tensor<DType, Any>?> =
         listOf(reluGrad(upstream, inputs[0], output))
 

skainet-compile/skainet-compile-dag/src/commonTest/kotlin/sk/ainet/compile/graph/ComputeGraphExecutorTest.kt

@@ -164,6 +164,7 @@ private class TestTensorOps : TensorOps {
     override fun <T : DType, V> rdivScalar(a: Number, b: Tensor<T, V>): Tensor<T, V> = b
     override fun <T : DType, V> matmul(a: Tensor<T, V>, b: Tensor<T, V>): Tensor<T, V> = a
     override fun <T : DType, V> transpose(tensor: Tensor<T, V>): Tensor<T, V> = tensor
+    override fun <T : DType, V> permute(tensor: Tensor<T, V>, axes: IntArray): Tensor<T, V> = tensor
     override fun <T : DType, V> relu(tensor: Tensor<T, V>): Tensor<T, V> = tensor
     override fun <T : DType, V> leakyRelu(tensor: Tensor<T, V>, negativeSlope: Float): Tensor<T, V> = tensor
     override fun <T : DType, V> elu(tensor: Tensor<T, V>, alpha: Float): Tensor<T, V> = tensor

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 <T : DType, V> transpose(tensor: Tensor<T, V>): Tensor<T, V>
 
+    /**
+     * 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 <T : DType, V> permute(tensor: Tensor<T, V>, axes: IntArray): Tensor<T, V>
+
     // Convolutional operations
     @Diff
     public fun <T : DType, V> conv1d(

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 <T : DType, V> permute(tensor: Tensor<T, V>, axes: IntArray): Tensor<T, V> {
+        validatePermuteAxes(tensor.shape, axes)
+        val resultShape = calculatePermuteShape(tensor.shape, axes)
+        val resultData = dataFactory.zeros<T, V>(resultShape, tensor.dtype)
+        return VoidOpsTensor(resultData, tensor.dtype)
+    }
+
     override fun <T : DType, V> conv1d(
         input: Tensor<T, V>,
         weight: Tensor<T, V>,
@@ -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")