diff --git a/skainet-backends/skainet-backend-cpu/src/commonTest/kotlin/sk/ainet/sk/ainet/exec/tensor/ops/LstmTest.kt b/skainet-backends/skainet-backend-cpu/src/commonTest/kotlin/sk/ainet/sk/ainet/exec/tensor/ops/LstmTest.kt new file mode 100644 index 00000000..cad43370 --- /dev/null +++ b/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( + 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(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(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(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(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(Shape(1, 1), FP32::class, floatArrayOf(0f)), + ctx.fromFloatArray(Shape(1, 1), FP32::class, floatArrayOf(0.8f)), // pre-charged cell + ) + val (out, next) = lstm.step(ctx.fromFloatArray(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) + } +} diff --git a/skainet-lang/skainet-lang-core/api/jvm/skainet-lang-core.api b/skainet-lang/skainet-lang-core/api/jvm/skainet-lang-core.api index 68950992..4c7cbad8 100644 --- a/skainet-lang/skainet-lang-core/api/jvm/skainet-lang-core.api +++ b/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 (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 (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 (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 (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 (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 (Lkotlin/Pair;)V public fun (Lkotlin/Pair;Lkotlin/Pair;)V diff --git a/skainet-lang/skainet-lang-core/src/commonMain/kotlin/sk/ainet/lang/nn/Lstm.kt b/skainet-lang/skainet-lang-core/src/commonMain/kotlin/sk/ainet/lang/nn/Lstm.kt new file mode 100644 index 00000000..47386c91 --- /dev/null +++ b/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( + public val h: Tensor, + public val c: Tensor, +) { + 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 @kotlin.jvm.JvmOverloads constructor( + public val inputSize: Int, + public val hiddenSize: Int, + override val name: String = "Lstm", + initWeightIh: Tensor, + initWeightHh: Tensor, + initBiasIh: Tensor, + initBiasHh: Tensor, + public val trainable: Boolean = true, +) : Module(), ModuleParameters { + + 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, 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, 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> = listOf(pWeightIh, pWeightHh, pBiasIh, pBiasHh) + + override val modules: List> + get() = emptyList() + + /** Zero initial state `(h0, c0)` for a batch of [batch] rows. */ + public fun initialState(batch: Int, ctx: ExecutionContext, dtype: kotlin.reflect.KClass): LstmState = + 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, state: LstmState, ctx: ExecutionContext): Pair, LstmState> { + 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, ctx: ExecutionContext): Tensor { + 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( + ctx.zeros(Shape(batch, hiddenSize), input.dtype), + ctx.zeros(Shape(batch, hiddenSize), input.dtype), + ) + + val outputs = ArrayList>(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] + } +}