jax-ml/jax · error · NotImplementedError

Dropout not supported in LSTM reference because we cannot de

Error message

Dropout not supported in LSTM reference because we cannot determine CUDNN dropout mask.

What it means

lstm_ref is a bit-exact reference for the cuDNN LSTM and cannot reproduce cuDNN's internal dropout RNG mask, so any nonzero dropout raises NotImplementedError.

Source

Thrown at jax/experimental/rnn.py:318

      precision=precision)
  return y, h_n, c_n


@jax.jit(static_argnums=(8, 9, 10, 11, 12))
def lstm_ref(x: Array, h_0: Array, c_0: Array, W_ih: dict[int, Array],
             W_hh: dict[int, Array], b_ih: dict[int, Array],
             b_hh: dict[int, Array], seq_lengths: Array, input_size: int,
             hidden_size: int, num_layers: int, dropout: float,
             bidirectional: bool) -> tuple[Array, Array, Array]:
  """Reference implementation of LSTM.

  See https://pytorch.org/docs/stable/generated/torch.nn.LSTM.html#lstm
  https://docs.nvidia.com/deeplearning/cudnn/api/index.html#cudnnRNNMode_t
  """
  if seq_lengths.dtype != jnp.dtype("int32"):
    raise NotImplementedError("`seq_lengths` can only be int32.")
  if dropout != 0.0:
    raise NotImplementedError(
        'Dropout not supported in LSTM reference because we cannot determine CUDNN dropout mask.'
    )

  # TODO(zhangqiaorjc): Handle ragged seq_lengths.
  # batch_size, max_seq_length = x.shape[0], x.shape[1]
  # assert seq_lengths.shape == (batch_size,)
  # for i in range(batch_size):
  #   if int(seq_lengths[i]) != max_seq_length:
  #     raise NotImplementedError('Does not yet support ragged sequences.')

  def lstm_cell(carry, x, *, W_ih, W_hh, b_ih, b_hh):
    h, c = carry
    W_ii, W_if, W_ig, W_io = jnp.split(W_ih, 4, axis=0)
    W_hi, W_hf, W_hg, W_ho = jnp.split(W_hh, 4, axis=0)
    b_ii, b_if, b_ig, b_io = jnp.split(b_ih, 4, axis=0)
    b_hi, b_hf, b_hg, b_ho = jnp.split(b_hh, 4, axis=0)
    i = sigmoid(x @ W_ii.T + b_ii[None] + h @ W_hi.T + b_hi[None])
    f = sigmoid(x @ W_if.T + b_if[None] + h @ W_hf.T + b_hf[None])

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

1. Set dropout=0.0 for reference-implementation comparisons
2. Compare against cuDNN with dropout disabled in both paths
3. Use the real cuDNN path (lstm, not lstm_ref) when dropout must be exercised

Example fix

# before
lstm_ref(x, h0, c0, w, sl, ..., dropout=0.1, ...)
# after
lstm_ref(x, h0, c0, w, sl, ..., dropout=0.0, ...)

assert dropout == 0.0, 'lstm_ref requires dropout=0'

When it happens

Trigger: Calling jax.experimental.rnn.lstm_ref with dropout != 0.0 while trying to verify gradients/outputs against the cuDNN implementation.

Common situations: Writing tests that compare the custom-VJL LSTM against the reference with the same dropout flags as training config.

Related errors

• bfloat16 support not implemented for LSTM
• Unexpected precision specifier value {precision}
• `seq_lengths` can only be int32.
• array ref with memory space only works inside of a `jit`.
• pinned array ref only works inside of a `jit`.

AI-assisted analysis of jax-ml/jax@1e1c6a8fc0 (2026-08-27). Data as JSON: /api/errors/272d4037e898b35c. Report an issue: GitHub.