jax-ml/jax · error · NotImplementedError

`seq_lengths` can only be int32.

Error message

`seq_lengths` can only be int32.

What it means

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.

Source

Thrown at jax/experimental/rnn.py:316

      dropout=dropout,
      bidirectional=bidirectional,
      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)

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

1. Cast seq_lengths to jnp.int32 before calling lstm_ref
2. Disable jax_enable_x64 if the script doesn't need 64-bit indexing

Example fix

# before
lstm_ref(x, h0, c0, w, seq_lengths_int64, ...)
# after
lstm_ref(x, h0, c0, w, seq_lengths.astype(jnp.int32), ...)

seq_lengths = jnp.asarray(seq_lengths, dtype=jnp.int32)

def is_int32_lengths(a) -> bool:
    return a.dtype == jnp.dtype('int32')

When it happens

Trigger: 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.

Common situations: 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.

Related errors

• bfloat16 support not implemented for LSTM
• Unexpected precision specifier value {precision}
• Dropout not supported in LSTM reference because we cannot de
• primal and tangent arguments to jax.jvp do not match; dtypes
• unexpected JAX type (e.g. shape/dtype) for gradient ref pass

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