Skip to content

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

Merged
michalharakal merged 1 commit into
developfrom
feature/lstm-layer
Jul 14, 2026
Merged

feat(nn): Lstm layer with explicit-state step() API#824
michalharakal merged 1 commit into
developfrom
feature/lstm-layer

Conversation

@michalharakal

Copy link
Copy Markdown
Contributor

Implements the DARC proposal #823 — C (Code) phase.

What

Lstm<T, V> in skainet-lang-core, mirroring Gru's unroll-at-trace-time design, built from
existing primitives only (matmul/narrow/sigmoid/tanh/multiply/reshape/unsqueeze/concat) — no
new TensorOps op, traces to StableHLO with no dedicated converter.

  • Gate order i, f, g, o and dual biases match torch.nn.LSTM; weights load after transposing
    to the matmul-ready [in, 4H] / [H, 4H] orientation (same convention as Gru).
  • New idiom: LstmState(h, c) + step(xt, state, ctx) — explicit caller-owned recurrent
    state, required by transducer prediction networks (RNN-T/TDT, e.g. Parakeet) and exactly the
    shape that lowers to a fixed-shape single-step StableHLO graph (state as graph I/O).
  • initialState(batch, ctx, dtype) helper.

Tests

LstmTest (skainet-backend-cpu, mirrors GruTest):

  • scalar-reference parity (all four gates + cell feedback), 1e-5
  • step() == unrolled forward() equivalence (the transducer usage pattern)
  • gate-order sanity: forget-gate cell persistence
  • output shapes

apiDump synced (jvm), apiCheck green.

Motivation / consumers

  • Parakeet FastConformer-TDT (SKaiNET-parakeet): 2-layer LSTM prediction network — blocking
  • Silero VAD port (SL2610 demo, PLAN 4.1) — the last runtime-Python piece

Closes #823

🤖 Generated with Claude Code

Single-layer, batch-first LSTM mirroring Gru's unroll-at-trace-time design,
built from existing primitives only (matmul/narrow/sigmoid/tanh/multiply) —
no new TensorOps op, traces to StableHLO with no dedicated converter.

- gate order i,f,g,o and dual biases match torch.nn.LSTM (weights load after
  transpose to the matmul-ready [in,4H]/[H,4H] orientation, like Gru)
- NEW: LstmState(h, c) + step(xt, state, ctx) — explicit caller-owned state,
  required by transducer prediction networks (RNN-T/TDT, e.g. Parakeet) and
  exactly the shape that lowers to a fixed-shape single-step StableHLO graph
- initialState(batch, ctx, dtype) helper
- LstmTest: scalar-reference parity (1e-5), step()==sequence equivalence,
  gate-order sanity (forget-gate cell persistence), output shapes
- api dump synced (jvm)

Unblocks: Parakeet-TDT prediction network (SKaiNET-parakeet P1/G2) and the
Silero VAD port (PLAN 4.1).

Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
@github-actions

Copy link
Copy Markdown

📖 Documentation Preview

The documentation has been built successfully for this PR.

Generated Files:

  • Operator documentation: docs/modules/operators/_generated_/
  • JSON schema output: operators.json

Artifacts:

  • Download the documentation-preview-824 artifact to view the complete documentation locally.

This comment will be updated automatically when the PR is updated.

@michalharakal
michalharakal merged commit 25a02b5 into develop Jul 14, 2026
14 checks passed
@michalharakal
michalharakal deleted the feature/lstm-layer branch July 14, 2026 20:00
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

[Feature]: LSTM layer (Lstm + explicit-state step API) in skainet-lang-core

1 participant