feat(nn): Lstm layer with explicit-state step() API

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

@@ -0,0 +1,136 @@
+package sk.ainet.sk.ainet.exec.tensor.ops
+
+import kotlin.math.exp
+import kotlin.math.tanh
+import kotlin.test.Test
+import kotlin.test.assertEquals
+import sk.ainet.context.DirectCpuExecutionContext
+import sk.ainet.lang.nn.Lstm
+import sk.ainet.lang.nn.LstmState
+import sk.ainet.lang.tensor.Shape
+import sk.ainet.lang.types.FP32
+
+/**
+ * Eager forward correctness for the LSTM layer, checked against an independent scalar
+ * reference (raw FloatArray loops) computing the same PyTorch i,f,g,o recurrence.
+ * Exercises all four gates, the cell-state feedback, and the step()/sequence equivalence
+ * (the transducer prediction-network usage pattern).
+ */
+class LstmTest {
+    private val ctx = DirectCpuExecutionContext()
+
+    private fun sigmoid(x: Float): Float = (1.0 / (1.0 + exp(-x.toDouble()))).toFloat()
+
+    // Independent reference: input [S*D] (batch=1), weights matmul-ready (row-major).
+    private fun lstmRef(
+        x: FloatArray, seq: Int, d: Int, h: Int,
+        wIh: FloatArray, wHh: FloatArray, bIh: FloatArray, bHh: FloatArray,
+    ): FloatArray {
+        val g4 = 4 * h
+        val hidden = FloatArray(h)
+        val cell = FloatArray(h)
+        val out = FloatArray(seq * h)
+        for (t in 0 until seq) {
+            val gx = FloatArray(g4) { k -> bIh[k] + (0 until d).sumOf { i -> (x[t * d + i] * wIh[i * g4 + k]).toDouble() }.toFloat() }
+            val gh = FloatArray(g4) { k -> bHh[k] + (0 until h).sumOf { j -> (hidden[j] * wHh[j * g4 + k]).toDouble() }.toFloat() }
+            for (j in 0 until h) {
+                val i = sigmoid(gx[j] + gh[j])                                   // input gate
+                val f = sigmoid(gx[h + j] + gh[h + j])                           // forget gate
+                val g = tanh((gx[2 * h + j] + gh[2 * h + j]).toDouble()).toFloat() // cell candidate
+                val o = sigmoid(gx[3 * h + j] + gh[3 * h + j])                   // output gate
+                cell[j] = f * cell[j] + i * g
+                hidden[j] = o * tanh(cell[j].toDouble()).toFloat()
+            }
+            for (j in 0 until h) out[t * h + j] = hidden[j]
+        }
+        return out
+    }
+
+    private fun makeLstm(d: Int, h: Int, wIh: FloatArray, wHh: FloatArray, bIh: FloatArray, bHh: FloatArray) =
+        Lstm<FP32, Float>(
+            inputSize = d, hiddenSize = h, name = "lstm",
+            initWeightIh = ctx.fromFloatArray(Shape(d, 4 * h), FP32::class, wIh),
+            initWeightHh = ctx.fromFloatArray(Shape(h, 4 * h), FP32::class, wHh),
+            initBiasIh = ctx.fromFloatArray(Shape(4 * h), FP32::class, bIh),
+            initBiasHh = ctx.fromFloatArray(Shape(4 * h), FP32::class, bHh),
+        )
+
+    @Test
+    fun lstm_forward_matches_reference() {
+        val batch = 1; val seq = 3; val d = 2; val h = 2; val g4 = 4 * h
+        // deterministic small weights/inputs (kept in sigmoid/tanh's sensitive range)
+        val x = FloatArray(seq * d) { ((it % 5) - 2) * 0.3f }
+        val wIh = FloatArray(d * g4) { ((it % 7) - 3) * 0.1f }
+        val wHh = FloatArray(h * g4) { ((it % 5) - 2) * 0.15f }
+        val bIh = FloatArray(g4) { ((it % 3) - 1) * 0.2f }
+        val bHh = FloatArray(g4) { ((it % 4) - 2) * 0.1f }
+
+        val lstm = makeLstm(d, h, wIh, wHh, bIh, bHh)
+        val input = ctx.fromFloatArray<FP32, Float>(Shape(batch, seq, d), FP32::class, x)
+        val out = lstm.forward(input, ctx)
+        assertEquals(Shape(batch, seq, h), out.shape)
+
+        val expected = lstmRef(x, seq, d, h, wIh, wHh, bIh, bHh)
+        for (t in 0 until seq) for (j in 0 until h) {
+            assertEquals(expected[t * h + j], out.data[0, t, j], 1e-5f)
+        }
+    }
+
+    @Test
+    fun lstm_step_equals_unrolled_forward() {
+        // The transducer usage: repeated step() with caller-owned state must reproduce the
+        // sequence forward exactly.
+        val batch = 1; val seq = 4; val d = 3; val h = 2; val g4 = 4 * h
+        val x = FloatArray(seq * d) { ((it % 6) - 3) * 0.25f }
+        val wIh = FloatArray(d * g4) { ((it % 9) - 4) * 0.08f }
+        val wHh = FloatArray(h * g4) { ((it % 7) - 3) * 0.12f }
+        val bIh = FloatArray(g4) { ((it % 5) - 2) * 0.1f }
+        val bHh = FloatArray(g4) { ((it % 3) - 1) * 0.15f }
+
+        val lstm = makeLstm(d, h, wIh, wHh, bIh, bHh)
+        val seqOut = lstm.forward(ctx.fromFloatArray<FP32, Float>(Shape(batch, seq, d), FP32::class, x), ctx)
+
+        var state = lstm.initialState(batch, ctx, FP32::class)
+        for (t in 0 until seq) {
+            val xt = ctx.fromFloatArray<FP32, Float>(Shape(batch, d), FP32::class, FloatArray(d) { x[t * d + it] })
+            val (out, next) = lstm.step(xt, state, ctx)
+            state = next
+            for (j in 0 until h) {
+                assertEquals(seqOut.data[0, t, j], out.data[0, j], 1e-6f)
+            }
+        }
+    }
+
+    @Test
+    fun lstm_output_shape_is_batch_seq_hidden() {
+        val batch = 2; val seq = 3; val d = 4; val h = 5; val g4 = 4 * h
+        val lstm = makeLstm(
+            d, h,
+            FloatArray(d * g4) { (it % 9 - 4) * 0.05f },
+            FloatArray(h * g4) { (it % 7 - 3) * 0.05f },
+            FloatArray(g4) { 0f },
+            FloatArray(g4) { 0f },
+        )
+        val input = ctx.fromFloatArray<FP32, Float>(Shape(batch, seq, d), FP32::class, FloatArray(batch * seq * d) { (it % 11 - 5) * 0.1f })
+        val out = lstm.forward(input, ctx)
+        assertEquals(Shape(batch, seq, h), out.shape)
+    }
+
+    @Test
+    fun forget_gate_controls_cell_memory() {
+        // With strongly positive forget-gate bias and zero input gate, the cell must persist;
+        // sanity-check the gate ORDER (i,f,g,o) is honored: a big b_ih in the f-slot must not
+        // leak into i/g/o behavior.
+        val d = 1; val h = 1; val g4 = 4
+        val bIh = floatArrayOf(-20f, 20f, 0f, 0f)   // i≈0, f≈1, g=0, o≈0.5
+        val lstm = makeLstm(d, h, FloatArray(d * g4), FloatArray(h * g4), bIh, FloatArray(g4))
+        val state0 = LstmState(
+            ctx.fromFloatArray<FP32, Float>(Shape(1, 1), FP32::class, floatArrayOf(0f)),
+            ctx.fromFloatArray<FP32, Float>(Shape(1, 1), FP32::class, floatArrayOf(0.8f)),  // pre-charged cell
+        )
+        val (out, next) = lstm.step(ctx.fromFloatArray<FP32, Float>(Shape(1, 1), FP32::class, floatArrayOf(1f)), state0, ctx)
+        // c' ≈ 1.0*0.8 + 0*0 = 0.8 ; h' ≈ 0.5 * tanh(0.8)
+        assertEquals(0.8f, next.c.data[0, 0], 1e-4f)
+        assertEquals(0.5f * tanh(0.8), out.data[0, 0].toDouble().toFloat().toDouble(), 1e-4)
+    }
+}

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

@@ -1000,6 +1000,27 @@ public final class sk/ainet/lang/nn/Linear : sk/ainet/lang/nn/Module, sk/ainet/l
 	public final fun getTrainable ()Z
 }
 
+public final class sk/ainet/lang/nn/Lstm : sk/ainet/lang/nn/Module, sk/ainet/lang/nn/topology/ModuleParameters {
+	public fun <init> (IILjava/lang/String;Lsk/ainet/lang/tensor/Tensor;Lsk/ainet/lang/tensor/Tensor;Lsk/ainet/lang/tensor/Tensor;Lsk/ainet/lang/tensor/Tensor;)V
+	public fun <init> (IILjava/lang/String;Lsk/ainet/lang/tensor/Tensor;Lsk/ainet/lang/tensor/Tensor;Lsk/ainet/lang/tensor/Tensor;Lsk/ainet/lang/tensor/Tensor;Z)V
+	public synthetic fun <init> (IILjava/lang/String;Lsk/ainet/lang/tensor/Tensor;Lsk/ainet/lang/tensor/Tensor;Lsk/ainet/lang/tensor/Tensor;Lsk/ainet/lang/tensor/Tensor;ZILkotlin/jvm/internal/DefaultConstructorMarker;)V
+	public fun <init> (IILsk/ainet/lang/tensor/Tensor;Lsk/ainet/lang/tensor/Tensor;Lsk/ainet/lang/tensor/Tensor;Lsk/ainet/lang/tensor/Tensor;)V
+	public final fun getHiddenSize ()I
+	public final fun getInputSize ()I
+	public fun getModules ()Ljava/util/List;
+	public fun getName ()Ljava/lang/String;
+	public fun getParams ()Ljava/util/List;
+	public final fun getTrainable ()Z
+	public final fun initialState (ILsk/ainet/context/ExecutionContext;Lkotlin/reflect/KClass;)Lsk/ainet/lang/nn/LstmState;
+	public final fun step (Lsk/ainet/lang/tensor/Tensor;Lsk/ainet/lang/nn/LstmState;Lsk/ainet/context/ExecutionContext;)Lkotlin/Pair;
+}
+
+public final class sk/ainet/lang/nn/LstmState {
+	public fun <init> (Lsk/ainet/lang/tensor/Tensor;Lsk/ainet/lang/tensor/Tensor;)V
+	public final fun getC ()Lsk/ainet/lang/tensor/Tensor;
+	public final fun getH ()Lsk/ainet/lang/tensor/Tensor;
+}
+
 public final class sk/ainet/lang/nn/MaxPool2d : sk/ainet/lang/nn/Module {
 	public fun <init> (Lkotlin/Pair;)V
 	public fun <init> (Lkotlin/Pair;Lkotlin/Pair;)V

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

@@ -0,0 +1,160 @@
+package sk.ainet.lang.nn
+
+import sk.ainet.context.ExecutionContext
+import sk.ainet.lang.tensor.Shape
+import sk.ainet.lang.tensor.Tensor
+import sk.ainet.lang.types.DType
+import sk.ainet.lang.nn.topology.ModuleParameter
+import sk.ainet.lang.nn.topology.ModuleParameters
+
+/**
+ * Recurrent state of an [Lstm]: hidden state [h] and cell state [c], both `[batch, hiddenSize]`.
+ *
+ * The state is EXPLICIT and caller-owned so that transducer-style decoders (RNN-T/TDT prediction
+ * networks) can advance it one token at a time via [Lstm.step] — and so a single step lowers to a
+ * fixed-shape StableHLO graph with the state as plain graph inputs/outputs (StableHLO has no loop
+ * construct; see [Gru] for the same unroll-at-trace-time rationale).
+ */
+public class LstmState<T : DType, V>(
+    public val h: Tensor<T, V>,
+    public val c: Tensor<T, V>,
+) {
+    init {
+        require(h.rank == 2 && c.rank == 2 && h.shape == c.shape) {
+            "LstmState: h and c must both be [batch, hidden], were ${h.shape} / ${c.shape}"
+        }
+    }
+}
+
+/**
+ * Single-layer, unidirectional, batch-first LSTM (long short-term memory).
+ *
+ * Input  `[batch, seq, inputSize]`  ->  output `[batch, seq, hiddenSize]` (all hidden states).
+ *
+ * Like [Gru], the recurrence is **unrolled over the (static) sequence length at trace time** and
+ * built entirely from existing primitive ops (matmul / add / sigmoid / tanh / multiply / narrow /
+ * concat), so it runs in the eager engine, is trainable through the standard autodiff tape, and
+ * traces to StableHLO with no dedicated converter. For step-wise decoding use [step] with an
+ * explicit [LstmState].
+ *
+ * Gate math matches `torch.nn.LSTM` (gate order **input, forget, cell/g, output**), so PyTorch
+ * weights load directly after transposing to the matmul-ready orientation used here:
+ *
+ *     i = sigmoid(x·W_ii + b_ii + h·W_hi + b_hi)
+ *     f = sigmoid(x·W_if + b_if + h·W_hf + b_hf)
+ *     g = tanh   (x·W_ig + b_ig + h·W_hg + b_hg)
+ *     o = sigmoid(x·W_io + b_io + h·W_ho + b_ho)
+ *     c' = f ⊙ c + i ⊙ g
+ *     h' = o ⊙ tanh(c')
+ *
+ * Weights are stored **matmul-ready** (input-major) — the four gates are concatenated on the
+ * trailing axis so a single matmul produces all four pre-activations:
+ *  - [weightIh] `[inputSize, 4*hiddenSize]`, [weightHh] `[hiddenSize, 4*hiddenSize]`
+ *  - [biasIh] / [biasHh] `[4*hiddenSize]`   (gate order i, f, g, o; PyTorch keeps both biases)
+ *
+ * Stacked LSTMs (e.g. a 2-layer transducer prediction network) compose single-layer cells so the
+ * per-layer state stays addressable in [step].
+ *
+ * @param inputSize  number of input features
+ * @param hiddenSize size of the hidden state
+ */
+public class Lstm<T : DType, V> @kotlin.jvm.JvmOverloads constructor(
+    public val inputSize: Int,
+    public val hiddenSize: Int,
+    override val name: String = "Lstm",
+    initWeightIh: Tensor<T, V>,
+    initWeightHh: Tensor<T, V>,
+    initBiasIh: Tensor<T, V>,
+    initBiasHh: Tensor<T, V>,
+    public val trainable: Boolean = true,
+) : Module<T, V>(), ModuleParameters<T, V> {
+
+    init {
+        require(inputSize > 0) { "Lstm($name): inputSize must be positive, was $inputSize" }
+        require(hiddenSize > 0) { "Lstm($name): hiddenSize must be positive, was $hiddenSize" }
+        val g = 4 * hiddenSize
+        fun check2d(t: Tensor<T, V>, rows: Int, cols: Int, what: String) {
+            val s = t.shape.dimensions
+            require(t.rank == 2 && s[0] == rows && s[1] == cols) {
+                "Lstm($name): $what shape must be [$rows, $cols], but was ${t.shape}"
+            }
+        }
+        check2d(initWeightIh, inputSize, g, "weightIh")
+        check2d(initWeightHh, hiddenSize, g, "weightHh")
+        fun check1d(t: Tensor<T, V>, len: Int, what: String) {
+            require(t.rank == 1 && t.shape.dimensions[0] == len) {
+                "Lstm($name): $what shape must be [$len], but was ${t.shape}"
+            }
+        }
+        check1d(initBiasIh, g, "biasIh")
+        check1d(initBiasHh, g, "biasHh")
+    }
+
+    private val pWeightIh = ModuleParameter.WeightParameter("$name.weight_ih", initWeightIh, trainable)
+    private val pWeightHh = ModuleParameter.WeightParameter("$name.weight_hh", initWeightHh, trainable)
+    private val pBiasIh = ModuleParameter.BiasParameter("$name.bias_ih", initBiasIh, trainable)
+    private val pBiasHh = ModuleParameter.BiasParameter("$name.bias_hh", initBiasHh, trainable)
+
+    override val params: List<ModuleParameter<T, V>> = listOf(pWeightIh, pWeightHh, pBiasIh, pBiasHh)
+
+    override val modules: List<Module<T, V>>
+        get() = emptyList()
+
+    /** Zero initial state `(h0, c0)` for a batch of [batch] rows. */
+    public fun initialState(batch: Int, ctx: ExecutionContext, dtype: kotlin.reflect.KClass<T>): LstmState<T, V> =
+        LstmState(
+            ctx.zeros(Shape(batch, hiddenSize), dtype),
+            ctx.zeros(Shape(batch, hiddenSize), dtype),
+        )
+
+    /**
+     * One recurrence step: `x_t [batch, inputSize]` + [state] -> `(h' [batch, hiddenSize], state')`.
+     * The returned hidden output IS `h'` (also carried inside the new state).
+     */
+    public fun step(xt: Tensor<T, V>, state: LstmState<T, V>, ctx: ExecutionContext): Pair<Tensor<T, V>, LstmState<T, V>> {
+        require(xt.rank == 2 && xt.shape[1] == inputSize) {
+            "Lstm($name): step input must be [batch, $inputSize], but was ${xt.shape}"
+        }
+        val ops = ctx.ops
+        val h = hiddenSize
+
+        // gate pre-activations: [batch, 4H], gate order i, f, g, o
+        val gx = ops.add(ops.matmul(xt, pWeightIh.value), pBiasIh.value)
+        val gh = ops.add(ops.matmul(state.h, pWeightHh.value), pBiasHh.value)
+        val gates = ops.add(gx, gh)
+
+        val i = ops.sigmoid(ops.narrow(gates, 1, 0, h))
+        val f = ops.sigmoid(ops.narrow(gates, 1, h, h))
+        val g = ops.tanh(ops.narrow(gates, 1, 2 * h, h))
+        val o = ops.sigmoid(ops.narrow(gates, 1, 3 * h, h))
+
+        val cNew = ops.add(ops.multiply(f, state.c), ops.multiply(i, g))
+        val hNew = ops.multiply(o, ops.tanh(cNew))
+        return hNew to LstmState(hNew, cNew)
+    }
+
+    override fun onForward(input: Tensor<T, V>, ctx: ExecutionContext): Tensor<T, V> {
+        require(input.rank == 3) {
+            "Lstm($name): input must be 3D [batch, seq, inputSize], but was ${input.shape}"
+        }
+        val ops = ctx.ops
+        val batch = input.shape[0]
+        val seq = input.shape[1]
+
+        // Initial state h0 = c0 = 0 (constant leaves in the trace).
+        var state = LstmState<T, V>(
+            ctx.zeros(Shape(batch, hiddenSize), input.dtype),
+            ctx.zeros(Shape(batch, hiddenSize), input.dtype),
+        )
+
+        val outputs = ArrayList<Tensor<T, V>>(seq)
+        for (t in 0 until seq) {
+            // x_t : [batch, inputSize]
+            val xt = ops.reshape(ops.narrow(input, 1, t, 1), Shape(batch, inputSize))
+            val (h, next) = step(xt, state, ctx)
+            state = next
+            outputs.add(ops.unsqueeze(h, 1)) // [batch, 1, H]
+        }
+        return ops.concat(outputs, 1) // [batch, seq, H]
+    }
+}