Implementation:Microsoft Onnxruntime CPU LSTM Forward
| Knowledge Sources | |
|---|---|
| Domains | Training, CPU_Kernels |
| Last Updated | 2026-02-10 04:00 GMT |
Overview
Concrete tool for running the LSTM forward pass during training on CPU in the ONNX Runtime training framework.
Description
This file implements the LSTMTraining kernel, which wraps the existing UniDirectionalLstm forward implementation with training mode enabled. It parses LSTM inputs and attributes via the lstm::LSTMInputs and lstm::LSTMOutputs utilities, then invokes the deep CPU LSTM computation. The kernel outputs all hidden states, the final hidden state, the final cell state, all cell states, and the iofc gate activations (input, output, forget, cell) needed for gradient computation in the backward pass. It is registered under the kMSDomain with opset version 1.
Usage
This kernel is invoked during the forward pass of LSTM training. The computed gate activations (iofc), all hidden states, and all cell states are stored as outputs and later consumed by the LSTMGrad kernel during backpropagation.
Code Reference
Source Location
- Repository: Microsoft_Onnxruntime
- File: orttraining/orttraining/training_ops/cpu/rnn/lstm.cc
- Lines: 1-64
Signature
template <typename T>
Status LSTMTraining<T>::Compute(OpKernelContext* context) const;
Import
#include "orttraining/orttraining/training_ops/cpu/rnn/lstm.h"
I/O Contract
Inputs
| Name | Type | Required | Description |
|---|---|---|---|
| X | Tensor(float) | Yes | Input sequence [seq_length, batch_size, input_size] |
| W | Tensor(float) | Yes | Weights [directions, 4*H, input_size] |
| R | Tensor(float) | Yes | Recurrence weights [directions, 4*H, H] |
| B | Tensor(float) | No | Bias [directions, 8*H] |
| sequence_lengths | Tensor(int) | No | Per-sequence lengths (not supported) |
| initial_h | Tensor(float) | No | Initial hidden state [directions, batch, H] |
| initial_c | Tensor(float) | No | Initial cell state [directions, batch, H] |
| P | Tensor(float) | No | Peephole weights [directions, 3*H] |
Outputs
| Name | Type | Description |
|---|---|---|
| HAll | Tensor(float) | All hidden states [seq_length, directions, batch, H] |
| Ht | Tensor(float) | Final hidden state [directions, batch, H] |
| Ct | Tensor(float) | Final cell state [directions, batch, H] |
| CAll | Tensor(float) | All cell states [seq_length, directions, batch, H] |
| IOFC | Tensor(float) | Gate activations [seq_length, directions, batch, 4*H] |
Usage Examples
ONNX_OPERATOR_TYPED_KERNEL_EX(
LSTMTraining, kMSDomain, 1, float, kCpuExecutionProvider,
(*KernelDefBuilder::Create())
.TypeConstraint("T", DataTypeImpl::GetTensorType<float>()),
LSTMTraining<float>);