Function lstmCell

Cell class for LSTM.

LSTMCell is distinct from the RNN subclass LSTM in that its apply method takes the input data of only a single time step and returns the cell's output at the time step, while LSTM takes the input data over a number of time steps. For example:

const cell = tf.layers.lstmCell({units: 2});
const input = tf.input({shape: [10]});
const output = cell.apply(input);

console.log(JSON.stringify(output.shape));
// [null, 10]: This is the cell's output at a single time step. The 1st
// dimension is the unknown batch size.
Copy

Instance(s) of LSTMCell can be used to construct RNN layers. The most typical use of this workflow is to combine a number of cells into a stacked RNN cell (i.e., StackedRNNCell internally) and use it to create an RNN. For example:

const cells = [
  tf.layers.lstmCell({units: 4}),
  tf.layers.lstmCell({units: 8}),
];
const rnn = tf.layers.rnn({cell: cells, returnSequences: true});

// Create an input with 10 time steps and a length-20 vector at each step.
const input = tf.input({shape: [10, 20]});
const output = rnn.apply(input);

console.log(JSON.stringify(output.shape));
// [null, 10, 8]: 1st dimension is unknown batch size; 2nd dimension is the
// same as the sequence length of `input`, due to `returnSequences`: `true`;
// 3rd dimension is the last `lstmCell`'s number of units.
Copy

To create an RNN consisting of only one LSTMCell, use the tf.layers.lstm.

Doc