Merge pull request #631 from SKaiNET-developers/feature/630-tanh-activation

Add tanh as a first-class TensorOps activation primitive (#630)

Author: michalharakal

skainet-backends/skainet-backend-cpu/src/commonMain/kotlin/sk/ainet/exec/tensor/ops/DefaultCpuOps.kt

@@ -2526,11 +2526,18 @@ public open class DefaultCpuOpsBase(protected val dataFactory: TensorDataFactory
         return newTensor(outData, tensor.dtype, tensor)
     }
 
+    @TensorOp()
+    @InProgress("cpu", owner = "team:cpu", issue = "task-ops.md#op-tanh")
     override fun <T : DType, V> tanh(tensor: Tensor<T, V>): Tensor<T, V> {
         val outData = dataFactory.init<T, V>(tensor.shape, tensor.dtype) { idx ->
-            val x = tensor.data.get(*idx) as Float
-            @Suppress("UNCHECKED_CAST")
-            kotlin.math.tanh(x).toFloat() as V
+            when (tensor.dtype) {
+                sk.ainet.lang.types.FP32::class, sk.ainet.lang.types.FP16::class -> {
+                    val x = tensor.data.get(*idx) as Float
+                    @Suppress("UNCHECKED_CAST")
+                    kotlin.math.tanh(x) as V
+                }
+                else -> throw IllegalArgumentException("Unsupported dtype for tanh: ${tensor.dtype}")
+            }
         }
         return newTensor(outData, tensor.dtype, tensor)
     }

skainet-backends/skainet-backend-cpu/src/commonTest/kotlin/sk/ainet/sk/ainet/exec/tensor/ops/DefaultCpuOpsActivationsTest.kt

@@ -27,6 +27,7 @@ class DefaultCpuOpsActivationsTest {
 
     private fun sigmoid(x: Float): Float = 1f / (1f + kotlin.math.exp(-x))
     private fun silu(x: Float): Float = x * sigmoid(x)
+    private fun tanh(x: Float): Float = kotlin.math.tanh(x)
 
     private fun assertAlmostEquals(expected: Float, actual: Float, eps: Float = 1e-5f, msg: String = "") {
         assertTrue(kotlin.math.abs(expected - actual) <= eps, msg.ifEmpty { "Expected $expected, got $actual" })
@@ -104,4 +105,50 @@ class DefaultCpuOpsActivationsTest {
         }
         assertTrue(threw, "Expected IllegalArgumentException for Int32 silu")
     }
+
+    @Test
+    fun tanh_fp32_basic_values() {
+        val input = fTensor(Shape(5), floatArrayOf(-2f, -1f, 0f, 1f, 2f))
+        val out = cpuOps.tanh(input)
+        assertEquals(Shape(5), out.shape)
+        assertEquals(FP32::class, out.dtype)
+        val expected = floatArrayOf(-2f, -1f, 0f, 1f, 2f).map { tanh(it) }
+        for (i in expected.indices) {
+            assertAlmostEquals(expected[i], out.data[i] as Float, 1e-6f, "tanh at $i mismatch")
+        }
+    }
+
+    @Test
+    fun tanh_fp32_matrix_shape_preserved() {
+        val input = fTensor(Shape(2, 3), floatArrayOf(
+            -1f, 0f, 1f,
+            2f, -2f, 0.5f
+        ))
+        val out = cpuOps.tanh(input)
+        assertEquals(Shape(2, 3), out.shape)
+        assertAlmostEquals(tanh(-1f), out.data[0, 0] as Float, 1e-6f)
+        assertAlmostEquals(tanh(0.5f), out.data[1, 2] as Float, 1e-6f)
+    }
+
+    @Test
+    fun tanh_fp32_saturates_at_extremes() {
+        // tanh(x) → 1 as x → +∞, → -1 as x → -∞. kotlin.math.tanh is numerically stable.
+        val input = fTensor(Shape(2), floatArrayOf(-100f, 100f))
+        val out = cpuOps.tanh(input)
+        assertAlmostEquals(-1f, out.data[0] as Float, 1e-6f)
+        assertAlmostEquals(1f, out.data[1] as Float, 1e-6f)
+    }
+
+    @Test
+    fun tanh_unsupported_dtype_int32_throws() {
+        val input = iTensor(Shape(2), intArrayOf(1, 2))
+        var threw = false
+        try {
+            cpuOps.tanh(input as sk.ainet.lang.tensor.Tensor<Int32, Int>)
+        } catch (e: IllegalArgumentException) {
+            threw = true
+            assertTrue(e.message?.contains("Unsupported dtype") == true)
+        }
+        assertTrue(threw, "Expected IllegalArgumentException for Int32 tanh")
+    }
 }

skainet-compile/skainet-compile-core/src/commonMain/kotlin/sk/ainet/tape/RecordingExecution.kt

@@ -136,6 +136,7 @@ private fun stableInputName(op: Operation, index: Int, total: Int): String = whe
     is ReluOperation<*, *> -> "input"
     is SoftmaxOperation<*, *> -> "input"
     is SigmoidOperation<*, *> -> "input"
+    is TanhOperation<*, *> -> "input"
     is SqueezeOperation<*, *> -> "input"
     is UnsqueezeOperation<*, *> -> "input"
     else -> if (total == 1) "input" else "input_$index"
@@ -418,6 +419,12 @@ internal class RecordingTensorOpsDecorator(private val base: TensorOps) : Tensor
         return out
     }
 
+    override fun <T : DType, V> tanh(tensor: Tensor<T, V>): Tensor<T, V> {
+        val out = base.tanh(tensor)
+        record(TanhOperation<T, V>(), listOf(tensor), listOf(out))
+        return out
+    }
+
     // --- Misc ---
     override fun <T : DType, V> squeeze(tensor: Tensor<T, V>, dim: Int?): Tensor<T, V> {
         val out = base.squeeze(tensor, dim)

skainet-compile/skainet-compile-dag/src/commonMain/kotlin/sk/ainet/lang/graph/DefaultExecutionTape.kt

@@ -175,6 +175,7 @@ public open class DefaultExecutionTape(
                     "matmul" -> listOf(ops.matmul(typedInputs[0], typedInputs[1]))
                     "relu" -> listOf(ops.relu(typedInputs[0]))
                     "sigmoid" -> listOf(ops.sigmoid(typedInputs[0]))
+                    "tanh" -> listOf(ops.tanh(typedInputs[0]))
                     "sum" -> listOf(ops.sum(typedInputs[0], params["dim"] as? Int))
                     "mean" -> listOf(ops.mean(typedInputs[0], params["dim"] as? Int))
                     "concat" -> listOf(ops.concat(typedInputs, params["dim"] as Int))
@@ -880,6 +881,13 @@ public class DefaultGradientTape(
         return listOf(grad)
     }
 
+    override fun tanhBackward(upstream: Tensor<DType, Any>, output: Tensor<DType, Any>, inputs: List<Tensor<DType, Any>>, attributes: Map<String, Any?>): List<Tensor<DType, Any>?> {
+        // d(tanh(x))/dx = 1 - tanh(x)^2 = 1 - output^2
+        val oneMinusSquare = output.ops.rsubScalar(1.0, output.ops.multiply(output, output))
+        val grad = upstream.ops.multiply(upstream, oneMinusSquare)
+        return listOf(grad)
+    }
+
     override fun siluBackward(upstream: Tensor<DType, Any>, output: Tensor<DType, Any>, inputs: List<Tensor<DType, Any>>, attributes: Map<String, Any?>): List<Tensor<DType, Any>?> {
         // silu(x) = x * sigmoid(x)
         // d(silu(x))/dx = sigmoid(x) + x * sigmoid(x) * (1 - sigmoid(x)) = sigmoid(x) + silu(x) * (1 - sigmoid(x))
@@ -1040,6 +1048,7 @@ public class DefaultGradientTape(
             "squeeze" -> BackwardOp(inputs, output) { upstream -> squeezeBackward(upstream, output, inputs, trace.attributes) }
             "unsqueeze" -> BackwardOp(inputs, output) { upstream -> unsqueezeBackward(upstream, output, inputs, trace.attributes) }
             "sigmoid" -> BackwardOp(inputs, output) { upstream -> sigmoidBackward(upstream, output, inputs, trace.attributes) }
+            "tanh" -> BackwardOp(inputs, output) { upstream -> tanhBackward(upstream, output, inputs, trace.attributes) }
             "silu" -> BackwardOp(inputs, output) { upstream -> siluBackward(upstream, output, inputs, trace.attributes) }
             "gelu" -> BackwardOp(inputs, output) { upstream -> geluBackward(upstream, output, inputs, trace.attributes) }
             "variance" -> BackwardOp(inputs, output) { upstream -> varianceBackward(upstream, output, inputs, trace.attributes) }

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

@@ -171,6 +171,7 @@ private class TestTensorOps : TensorOps {
     override fun <T : DType, V> softmax(tensor: Tensor<T, V>, dim: Int): Tensor<T, V> = tensor
     override fun <T : DType, V> logSoftmax(tensor: Tensor<T, V>, dim: Int): Tensor<T, V> = tensor
     override fun <T : DType, V> sigmoid(tensor: Tensor<T, V>): Tensor<T, V> = tensor
+    override fun <T : DType, V> tanh(tensor: Tensor<T, V>): Tensor<T, V> = tensor
     override fun <T : DType, V> silu(tensor: Tensor<T, V>): Tensor<T, V> = tensor
     override fun <T : DType, V> gelu(tensor: Tensor<T, V>): Tensor<T, V> = tensor
     override fun <T : DType, V> sum(tensor: Tensor<T, V>, dim: Int?): Tensor<T, V> = tensor

skainet-lang/skainet-lang-core/src/commonMain/kotlin/sk/ainet/lang/tensor/TensorExtensions.kt

@@ -73,6 +73,7 @@ public fun <T : DType, V> Tensor<T, V>.relu(): Tensor<T, V> = ops.relu(this)
 public fun <T : DType, V> Tensor<T, V>.leakyRelu(negativeSlope: Float = 0.01f): Tensor<T, V> = ops.leakyRelu(this, negativeSlope)
 public fun <T : DType, V> Tensor<T, V>.elu(alpha: Float = 1.0f): Tensor<T, V> = ops.elu(this, alpha)
 public fun <T : DType, V> Tensor<T, V>.sigmoid(): Tensor<T, V> = ops.sigmoid(this)
+public fun <T : DType, V> Tensor<T, V>.tanh(): Tensor<T, V> = ops.tanh(this)
 public fun <T : DType, V> Tensor<T, V>.silu(): Tensor<T, V> = ops.silu(this)
 public fun <T : DType, V> Tensor<T, V>.gelu(): Tensor<T, V> = ops.gelu(this)
 public fun <T : DType, V> Tensor<T, V>.exp(): Tensor<T, V> = ops.exp(this)

skainet-lang/skainet-lang-core/src/commonMain/kotlin/sk/ainet/lang/tensor/ops/TensorOperations.kt

@@ -813,6 +813,37 @@ public class SigmoidOperation<T : DType, V>(
     override fun clone(newParameters: Map<String, Any>): Operation = SigmoidOperation<T, V>(newParameters)
 }
 
+public class TanhOperation<T : DType, V>(
+    parameters: Map<String, Any> = emptyMap()
+) : BaseOperation("tanh", "activation", parameters) {
+
+    override fun <T2 : DType, V2> execute(inputs: List<Tensor<T2, V2>>): List<Tensor<T2, V2>> {
+        require(inputs.size == 1) { "Tanh operation requires exactly 1 input" }
+        throw UnsupportedOperationException("Direct execution not supported in graph mode")
+    }
+
+    override fun validateInputs(inputs: List<TensorSpec>): ValidationResult {
+        if (inputs.size != 1) {
+            return ValidationResult.Invalid(listOf("Tanh operation requires exactly 1 input, got ${inputs.size}"))
+        }
+        return ValidationResult.Valid
+    }
+
+    override fun inferOutputs(inputs: List<TensorSpec>): List<TensorSpec> {
+        require(inputs.size == 1) { "Tanh operation requires exactly 1 input" }
+        return listOf(
+            TensorSpec(
+                name = "tanh_output",
+                shape = inputs[0].shape,
+                dtype = inputs[0].dtype,
+                requiresGrad = inputs[0].requiresGrad
+            )
+        )
+    }
+
+    override fun clone(newParameters: Map<String, Any>): Operation = TanhOperation<T, V>(newParameters)
+}
+
 /**
  * Additional shape operations
  */

skainet-lang/skainet-lang-core/src/commonMain/kotlin/sk/ainet/lang/tensor/ops/TensorOps.kt

@@ -178,6 +178,9 @@ public interface TensorOps {
     public fun <T : DType, V> sigmoid(tensor: Tensor<T, V>): Tensor<T, V>
     @Diff
     @ActivationDsl
+    public fun <T : DType, V> tanh(tensor: Tensor<T, V>): Tensor<T, V>
+    @Diff
+    @ActivationDsl
     public fun <T : DType, V> silu(tensor: Tensor<T, V>): Tensor<T, V>
     @Diff
     @ActivationDsl
@@ -304,10 +307,6 @@ public interface TensorOps {
         throw NotImplementedError("cos not implemented by this TensorOps backend")
     }
 
-    public fun <T : DType, V> tanh(tensor: Tensor<T, V>): Tensor<T, V> {
-        throw NotImplementedError("tanh not implemented by this TensorOps backend")
-    }
-
     /**
      * Scaled dot-product attention.
      *