jax-ml/jax · error · NotImplementedError

bfloat16 support not implemented for LSTM

Error message

bfloat16 support not implemented for LSTM

What it means

The experimental cuDNN-backed LSTM in jax.experimental.rnn does not support bfloat16 computations. The precision mapping function treats lax.Precision.DEFAULT as bfloat16 mode and raises NotImplementedError because cuDNN LSTM lacks bf16 support on that path.

Source

Thrown at jax/experimental/rnn.py:265

  # the logic from canonicalize_precision that we require here boils down to:
  #
  #   if precision is None and config.jax_default_matmul_precision is not None:
  #     precision = Precision(config.jax_default_matmul_precision)
  #   else:
  #     precision = None
  #
  # but we prefer to still invoke it here for consistency
  precision = lax.canonicalize_precision(precision)
  if precision is None or not (isinstance(precision, tuple) and len(precision) == 2):
    return True
  # cuDNN allows only one precision specifier per RNN op
  match precision:
    case (lax.Precision.HIGHEST, _):
      return False
    case (lax.Precision.HIGH, _):
      return True
    case (lax.Precision.DEFAULT, _): # bfloat16
      raise NotImplementedError("bfloat16 support not implemented for LSTM")
    case _:
      raise ValueError(f"Unexpected precision specifier value {precision}")


@partial(custom_vjp, nondiff_argnums=(5, 6, 7, 8, 9, 10))
def lstm(x: Array, h_0: Array, c_0: Array, weights: Array, seq_lengths: Array,
         input_size: int, hidden_size: int, num_layers: int, dropout: float,
         bidirectional: bool, precision: lax.PrecisionLike = None) -> tuple[Array, Array, Array]:
  """LSTM via CuDNN or HIPDNN (not-yet-supported).

  Assume batch-first inputs.

  Arguments:
    x: (batch_size, max_seq_length, input_size)
    h_0: (num_directions * num_layers, batch_size, hidden_size)
    c_0: (num_directions * num_layers, batch_size, hidden_size)
    weights: (num_params,) where num_params = get_num_params_in_lstm(...)
    seq_lengths: (batch_size,)

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

1. Cast LSTM inputs/weights to float32 and pass precision=lax.Precision.HIGHEST to keep FP32 math
2. Use precision=lax.Precision.HIGH to allow TF32 (still FP32 storage) on Ampere+ GPUs
3. Fall back to a manual LSTM via lax.scan for bf16 training

Example fix

# before
y, h_n, c_n = lstm(x_bf16, h0, c0, w_bf16, seq_lens, ..., precision=lax.Precision.DEFAULT)
# after
y, h_n, c_n = lstm(x_bf16.astype(jnp.float32), h0, c0, w.astype(jnp.float32), seq_lens, ..., precision=lax.Precision.HIGHEST)

if precision is None or precision == lax.Precision.DEFAULT:
    x, h0, c0, w = jax.tree.map(lambda a: a.astype(jnp.float32), (x, h0, c0, w))
    precision = lax.Precision.HIGHEST

try:
    out = lstm(...)
except NotImplementedError:
    out = my_lax_scan_lstm(...)  # fp32 or manual fallback

When it happens

Trigger: Calling jax.experimental.rnn.lstm with precision=lax.Precision.DEFAULT (or a tuple containing it) while running in a context where DEFAULT resolves to bfloat16 (e.g. bf16 params on GPU with cuDNN backend).

Common situations: Running bf16 LSTM training on A100/H100 GPUs where the rest of the model is bf16; passing precision=None which falls into the DEFAULT branch on some versions.

Related errors

• cuDNN doesn't support right window: {r_window} when causal m
• Unexpected precision specifier value {precision}
• `seq_lengths` can only be int32.
• Dropout not supported in LSTM reference because we cannot de
• cuDNN not found.

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