feat(nn): implement stochastic masking in Dropout

skainet-backends/skainet-backend-cpu/src/commonTest/kotlin/sk/ainet/exec/nn/DropoutMaskingTest.kt

@@ -0,0 +1,62 @@
+package sk.ainet.exec.nn
+
+import kotlin.math.abs
+import kotlin.random.Random
+import kotlin.test.Test
+import kotlin.test.assertContentEquals
+import kotlin.test.assertTrue
+import sk.ainet.context.DirectCpuExecutionContext
+import sk.ainet.context.Phase
+import sk.ainet.lang.nn.layers.Dropout
+import sk.ainet.lang.tensor.Shape
+import sk.ainet.lang.tensor.Tensor
+import sk.ainet.lang.types.FP32
+
+/**
+ * Value-level dropout behavior on the real CPU backend (the lang-core default
+ * context uses the shape-only Void backend, so these assertions live here).
+ */
+class DropoutMaskingTest {
+
+    private val trainCtx = DirectCpuExecutionContext(phase = Phase.TRAIN)
+    private val evalCtx = DirectCpuExecutionContext(phase = Phase.EVAL)
+
+    private fun ones(n: Int, ctx: DirectCpuExecutionContext): Tensor<FP32, Float> =
+        ctx.fromFloatArray(Shape(n), FP32::class, FloatArray(n) { 1f })
+
+    @Test
+    fun training_phase_zeroes_or_scales_every_element() {
+        val layer = Dropout<FP32, Float>(p = 0.5f, random = Random(42))
+        val out = layer.forward(ones(64, trainCtx), trainCtx).data.copyToFloatArray()
+
+        assertTrue(out.all { it == 0f || it == 2f }, "inverted dropout must zero or scale by 1/(1-p)")
+        assertTrue(out.any { it == 0f }, "expected at least one dropped element")
+        assertTrue(out.any { it == 2f }, "expected at least one surviving element")
+    }
+
+    @Test
+    fun expected_activation_is_preserved() {
+        val layer = Dropout<FP32, Float>(p = 0.3f, random = Random(123))
+        val out = layer.forward(ones(10_000, trainCtx), trainCtx).data.copyToFloatArray()
+
+        val mean = out.average().toFloat()
+        assertTrue(abs(mean - 1f) < 0.05f, "inverted dropout should keep the mean ~1, was $mean")
+    }
+
+    @Test
+    fun seeded_random_is_reproducible() {
+        val x = trainCtx.fromFloatArray<FP32, Float>(Shape(16), FP32::class, FloatArray(16) { it.toFloat() })
+
+        val y1 = Dropout<FP32, Float>(p = 0.5f, random = Random(7)).forward(x, trainCtx)
+        val y2 = Dropout<FP32, Float>(p = 0.5f, random = Random(7)).forward(x, trainCtx)
+
+        assertContentEquals(y1.data.copyToFloatArray(), y2.data.copyToFloatArray())
+    }
+
+    @Test
+    fun inference_phase_is_exact_identity() {
+        val x = evalCtx.fromFloatArray<FP32, Float>(Shape(4), FP32::class, floatArrayOf(1f, 2f, 3f, 4f))
+        val y = Dropout<FP32, Float>(p = 0.5f, random = Random(1)).forward(x, evalCtx)
+        assertContentEquals(floatArrayOf(1f, 2f, 3f, 4f), y.data.copyToFloatArray())
+    }
+}

skainet-lang/skainet-lang-core/api/android/skainet-lang-core.api

@@ -1641,8 +1641,8 @@ public final class sk/ainet/lang/nn/hooks/TapeRecorder$Entry {
 
 public final class sk/ainet/lang/nn/layers/Dropout : sk/ainet/lang/nn/Module {
 	public fun <init> ()V
-	public fun <init> (FZLjava/lang/String;)V
-	public synthetic fun <init> (FZLjava/lang/String;ILkotlin/jvm/internal/DefaultConstructorMarker;)V
+	public fun <init> (FZLjava/lang/String;Lkotlin/random/Random;)V
+	public synthetic fun <init> (FZLjava/lang/String;Lkotlin/random/Random;ILkotlin/jvm/internal/DefaultConstructorMarker;)V
 	public fun forward (Lsk/ainet/lang/tensor/Tensor;Lsk/ainet/context/ExecutionContext;)Lsk/ainet/lang/tensor/Tensor;
 	public fun getModules ()Ljava/util/List;
 	public fun getName ()Ljava/lang/String;

skainet-lang/skainet-lang-core/api/jvm/skainet-lang-core.api

@@ -1930,8 +1930,8 @@ public final class sk/ainet/lang/nn/hooks/TapeRecorder$Entry {
 
 public final class sk/ainet/lang/nn/layers/Dropout : sk/ainet/lang/nn/Module {
 	public fun <init> ()V
-	public fun <init> (FZLjava/lang/String;)V
-	public synthetic fun <init> (FZLjava/lang/String;ILkotlin/jvm/internal/DefaultConstructorMarker;)V
+	public fun <init> (FZLjava/lang/String;Lkotlin/random/Random;)V
+	public synthetic fun <init> (FZLjava/lang/String;Lkotlin/random/Random;ILkotlin/jvm/internal/DefaultConstructorMarker;)V
 	public fun forward (Lsk/ainet/lang/tensor/Tensor;Lsk/ainet/context/ExecutionContext;)Lsk/ainet/lang/tensor/Tensor;
 	public fun getModules ()Ljava/util/List;
 	public fun getName ()Ljava/lang/String;

skainet-lang/skainet-lang-core/src/commonMain/kotlin/sk/ainet/lang/nn/layers/Dropout.kt

@@ -3,19 +3,37 @@ package sk.ainet.lang.nn.layers
 import sk.ainet.context.ExecutionContext
 import sk.ainet.lang.nn.Module
 import sk.ainet.lang.tensor.Tensor
+import sk.ainet.lang.tensor.data.TensorData
+import sk.ainet.lang.types.BF16
 import sk.ainet.lang.types.DType
+import sk.ainet.lang.types.FP16
+import sk.ainet.lang.types.FP32
+import sk.ainet.lang.types.FP64
+import kotlin.random.Random
 
 /**
- * Dropout layer that is aware of ExecutionContext phases.
+ * Dropout layer with inverted-dropout semantics, aware of ExecutionContext phases.
  *
- * Current implementation keeps identity semantics on both TRAIN and EVAL phases,
- * but exposes a context-aware forward that can later be extended to apply
- * stochastic masking when ctx.inTraining is true.
+ * During training ([ExecutionContext.inTraining] and [training] both true, `p > 0`),
+ * each element is zeroed with probability [p] and the surviving elements are
+ * scaled by `1 / (1 - p)` so the expected activation stays constant. In the
+ * evaluation phase (or with [training] disabled) the layer is the identity —
+ * no rescaling is needed at inference time.
+ *
+ * The mask is a constant tensor combined with an element-wise multiply, so
+ * gradients flow to the surviving inputs only — the standard dropout backward —
+ * without needing a dedicated autograd rule.
+ *
+ * @param p probability of zeroing an element, in `[0, 1)`
+ * @param training manual override; set to false to disable masking regardless of phase
+ * @param name name of the module
+ * @param random RNG for the mask — inject a seeded [Random] for reproducible runs
  */
 public class Dropout<T : DType, V>(
     public val p: Float = 0.5f,
     public var training: Boolean = true,
-    override val name: String = "Dropout"
+    override val name: String = "Dropout",
+    private val random: Random = Random.Default,
 ) : Module<T, V>() {
 
     init {
@@ -27,12 +45,31 @@ public class Dropout<T : DType, V>(
         get() = emptyList()
 
     /**
-     * Context-aware forward that can use ExecutionContext.phase. Hooks are dispatched if available.
+     * Context-aware forward: stochastic masking when ctx is in the training
+     * phase, identity otherwise. Hooks are dispatched if available.
      */
     override fun forward(input: Tensor<T, V>, ctx: ExecutionContext): Tensor<T, V> =
         sk.ainet.lang.nn.hooks.withForwardHooks(ctx, this, input) {
-            // Placeholder behavior: identity in both phases. When RNG and elementwise ops are available,
-            // implement: if (ctx.inTraining && p > 0f) then output = input * mask / (1 - p)
-            input
+            if (!ctx.inTraining || !training || p == 0f) {
+                input
+            } else {
+                val mask = ctx.fromData(bernoulliMask(input, ctx), input.dtype)
+                ctx.ops.multiply(input, mask)
+            }
         }
+
+    private fun bernoulliMask(input: Tensor<T, V>, ctx: ExecutionContext): TensorData<T, V> {
+        val scale = 1f / (1f - p)
+        return ctx.tensorDataFactory.init(input.shape, input.dtype) {
+            val keep = random.nextFloat() >= p
+            @Suppress("UNCHECKED_CAST")
+            when (input.dtype) {
+                FP32::class, FP16::class, BF16::class -> (if (keep) scale else 0f) as V
+                FP64::class -> (if (keep) scale.toDouble() else 0.0) as V
+                else -> throw UnsupportedOperationException(
+                    "Dropout($name): unsupported dtype ${input.dtype} — floating-point tensors only"
+                )
+            }
+        }
+    }
 }

skainet-lang/skainet-lang-core/src/commonTest/kotlin/sk/ainet/lang/nn/layers/DropoutPhaseTest.kt

@@ -17,13 +17,15 @@ class DropoutPhaseTest {
         val layer = Dropout<FP32, Float>(p = 0.3f)
         val y = eval(base) { ctx -> layer.forward(x, ctx) }
         assertEquals(x.shape, y.shape)
+        kotlin.test.assertContentEquals(floatArrayOf(1f, 2f, 3f, 4f), y.data.copyToFloatArray())
     }
 
     @Test
-    fun dropout_identity_in_train_phase_with_ctx() {
+    fun dropout_identity_in_train_phase_when_p_is_zero() {
         val x = base.fromFloatArray<FP32, Float>(Shape(2, 2), FP32::class, floatArrayOf(1f, 2f, 3f, 4f))
         val layer = Dropout<FP32, Float>(p = 0.0f)
         val y = train(base) { ctx -> layer.forward(x, ctx) }
         assertEquals(x.shape, y.shape)
+        kotlin.test.assertContentEquals(floatArrayOf(1f, 2f, 3f, 4f), y.data.copyToFloatArray())
     }
 }

skainet-lang/skainet-lang-core/src/commonTest/kotlin/sk/ainet/lang/nn/layers/DropoutTest.kt

@@ -1,12 +1,20 @@
 package sk.ainet.lang.nn.layers
 
+import kotlin.random.Random
 import kotlin.test.Test
+import kotlin.test.assertContentEquals
 import kotlin.test.assertEquals
+import sk.ainet.context.train
+import sk.ainet.lang.nn.DefaultNeuralNetworkExecutionContext
+import sk.ainet.lang.nn.NeuralNetworkExecutionContext
 import sk.ainet.lang.tensor.Shape
 import sk.ainet.lang.types.FP32
-import sk.ainet.lang.nn.NeuralNetworkExecutionContext
-import sk.ainet.lang.nn.DefaultNeuralNetworkExecutionContext
 
+/**
+ * Identity paths and shape behavior. The default context here uses the
+ * shape-only Void backend, so value-level masking behavior is covered in
+ * skainet-backend-cpu's DropoutMaskingTest.
+ */
 class DropoutTest {
 
     private val ctx: NeuralNetworkExecutionContext = DefaultNeuralNetworkExecutionContext()
@@ -15,15 +23,23 @@ class DropoutTest {
     fun p_zero_is_identity_in_training() {
         val x = ctx.fromFloatArray<FP32, Float>(Shape(3), FP32::class, floatArrayOf(1f, 2f, 3f))
         val layer = Dropout<FP32, Float>(p = 0f, training = true)
-        val y = layer.forward(x, ctx)
-        assertEquals(x.shape, y.shape)
+        val y = train(ctx) { trainCtx -> layer.forward(x, trainCtx) }
+        assertContentEquals(floatArrayOf(1f, 2f, 3f), y.data.copyToFloatArray())
     }
 
     @Test
     fun eval_mode_is_identity() {
         val x = ctx.fromFloatArray<FP32, Float>(Shape(2, 2), FP32::class, floatArrayOf(1f, 2f, 3f, 4f))
         val layer = Dropout<FP32, Float>(p = 0.5f, training = false)
         val y = layer.forward(x, ctx)
+        assertContentEquals(floatArrayOf(1f, 2f, 3f, 4f), y.data.copyToFloatArray())
+    }
+
+    @Test
+    fun masked_path_preserves_the_shape() {
+        val x = ctx.fromFloatArray<FP32, Float>(Shape(2, 3, 4), FP32::class, FloatArray(24) { 1f })
+        val layer = Dropout<FP32, Float>(p = 0.5f, random = Random(1))
+        val y = train(ctx) { trainCtx -> layer.forward(x, trainCtx) }
         assertEquals(x.shape, y.shape)
     }
 }