{"record":{"id":"2fd588f095c17235","repo":"jax-ml/jax","slug":"seq-lengths-can-only-be-int32","errorCode":null,"errorMessage":"`seq_lengths` can only be int32.","messagePattern":"`seq_lengths` can only be int32\\.","errorType":"exception","errorClass":"NotImplementedError","httpStatus":null,"severity":"error","filePath":"jax/experimental/rnn.py","lineNumber":316,"sourceCode":"      dropout=dropout,\n      bidirectional=bidirectional,\n      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)","sourceCodeStart":298,"sourceCodeEnd":334,"githubUrl":"https://github.com/jax-ml/jax/blob/1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb/jax/experimental/rnn.py#L298-L334","documentation":"The reference (pure-JAX) LSTM implementation only accepts int32 seq_lengths because the reference semantics mirror cuDNN's fixed int32 length arrays. Any other dtype (int64, uint32) raises NotImplementedError in lstm_ref.","triggerScenarios":"Calling jax.experimental.rnn.lstm_ref with seq_lengths as int64 (default when jax_enable_x64=True) or computed via numpy operations that yield int64.","commonSituations":"Enabling 64-bit mode globally; building seq_lengths with np.arange or sums that stay int64 on Linux; loading seq_lengths from a dataset stored as int64.","solutions":["Cast seq_lengths to jnp.int32 before calling lstm_ref","Disable jax_enable_x64 if the script doesn't need 64-bit indexing"],"exampleFix":"# before\nlstm_ref(x, h0, c0, w, seq_lengths_int64, ...)\n# after\nlstm_ref(x, h0, c0, w, seq_lengths.astype(jnp.int32), ...)","handlingStrategy":"validation","validationCode":"seq_lengths = jnp.asarray(seq_lengths, dtype=jnp.int32)","typeGuard":"def is_int32_lengths(a) -> bool:\n    return a.dtype == jnp.dtype('int32')","tryCatchPattern":null,"preventionTips":["Normalize dtypes at data-loader boundary"],"tags":["jax","lstm","rnn","dtype","int32"],"backgroundTag":"dtype-mismatch","analyzedSha":"1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb","analyzedAt":"2026-08-27T09:53:25.647Z","schemaVersion":2},"datasetVersion":"2026-08-27T13:17:12.746Z"}