{"record":{"id":"272d4037e898b35c","repo":"jax-ml/jax","slug":"dropout-not-supported-in-lstm-reference-because-we","errorCode":null,"errorMessage":"Dropout not supported in LSTM reference because we cannot determine CUDNN dropout mask.","messagePattern":"Dropout not supported in LSTM reference because we cannot determine CUDNN dropout mask\\.","errorType":"exception","errorClass":"NotImplementedError","httpStatus":null,"severity":"error","filePath":"jax/experimental/rnn.py","lineNumber":318,"sourceCode":"      precision=precision)\n  return y, h_n, c_n\n\n\n@jax.jit(static_argnums=(8, 9, 10, 11, 12))\ndef lstm_ref(x: Array, h_0: Array, c_0: Array, W_ih: dict[int, Array],\n             W_hh: dict[int, Array], b_ih: dict[int, Array],\n             b_hh: dict[int, Array], seq_lengths: Array, input_size: int,\n             hidden_size: int, num_layers: int, dropout: float,\n             bidirectional: bool) -> tuple[Array, Array, Array]:\n  \"\"\"Reference implementation of LSTM.\n\n  See https://pytorch.org/docs/stable/generated/torch.nn.LSTM.html#lstm\n  https://docs.nvidia.com/deeplearning/cudnn/api/index.html#cudnnRNNMode_t\n  \"\"\"\n  if seq_lengths.dtype != jnp.dtype(\"int32\"):\n    raise NotImplementedError(\"`seq_lengths` can only be int32.\")\n  if dropout != 0.0:\n    raise NotImplementedError(\n        'Dropout not supported in LSTM reference because we cannot determine CUDNN dropout mask.'\n    )\n\n  # TODO(zhangqiaorjc): Handle ragged seq_lengths.\n  # batch_size, max_seq_length = x.shape[0], x.shape[1]\n  # assert seq_lengths.shape == (batch_size,)\n  # for i in range(batch_size):\n  #   if int(seq_lengths[i]) != max_seq_length:\n  #     raise NotImplementedError('Does not yet support ragged sequences.')\n\n  def lstm_cell(carry, x, *, W_ih, W_hh, b_ih, b_hh):\n    h, c = carry\n    W_ii, W_if, W_ig, W_io = jnp.split(W_ih, 4, axis=0)\n    W_hi, W_hf, W_hg, W_ho = jnp.split(W_hh, 4, axis=0)\n    b_ii, b_if, b_ig, b_io = jnp.split(b_ih, 4, axis=0)\n    b_hi, b_hf, b_hg, b_ho = jnp.split(b_hh, 4, axis=0)\n    i = sigmoid(x @ W_ii.T + b_ii[None] + h @ W_hi.T + b_hi[None])\n    f = sigmoid(x @ W_if.T + b_if[None] + h @ W_hf.T + b_hf[None])","sourceCodeStart":300,"sourceCodeEnd":336,"githubUrl":"https://github.com/jax-ml/jax/blob/1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb/jax/experimental/rnn.py#L300-L336","documentation":"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.","triggerScenarios":"Calling jax.experimental.rnn.lstm_ref with dropout != 0.0 while trying to verify gradients/outputs against the cuDNN implementation.","commonSituations":"Writing tests that compare the custom-VJL LSTM against the reference with the same dropout flags as training config.","solutions":["Set dropout=0.0 for reference-implementation comparisons","Compare against cuDNN with dropout disabled in both paths","Use the real cuDNN path (lstm, not lstm_ref) when dropout must be exercised"],"exampleFix":"# before\nlstm_ref(x, h0, c0, w, sl, ..., dropout=0.1, ...)\n# after\nlstm_ref(x, h0, c0, w, sl, ..., dropout=0.0, ...)","handlingStrategy":"validation","validationCode":"assert dropout == 0.0, 'lstm_ref requires dropout=0'","typeGuard":null,"tryCatchPattern":null,"preventionTips":["Treat lstm_ref as a dropout=0 verification tool only"],"tags":["jax","lstm","rnn","dropout","not-implemented"],"backgroundTag":"unsupported-feature","analyzedSha":"1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb","analyzedAt":"2026-08-27T09:53:25.647Z","schemaVersion":2},"datasetVersion":"2026-08-27T13:17:12.746Z"}