{"record":{"id":"d3d04f45ec560043","repo":"jax-ml/jax","slug":"bfloat16-support-not-implemented-for-lstm","errorCode":null,"errorMessage":"bfloat16 support not implemented for LSTM","messagePattern":"bfloat16 support not implemented for LSTM","errorType":"exception","errorClass":"NotImplementedError","httpStatus":null,"severity":"error","filePath":"jax/experimental/rnn.py","lineNumber":265,"sourceCode":"  # the logic from canonicalize_precision that we require here boils down to:\n  #\n  #   if precision is None and config.jax_default_matmul_precision is not None:\n  #     precision = Precision(config.jax_default_matmul_precision)\n  #   else:\n  #     precision = None\n  #\n  # but we prefer to still invoke it here for consistency\n  precision = lax.canonicalize_precision(precision)\n  if precision is None or not (isinstance(precision, tuple) and len(precision) == 2):\n    return True\n  # cuDNN allows only one precision specifier per RNN op\n  match precision:\n    case (lax.Precision.HIGHEST, _):\n      return False\n    case (lax.Precision.HIGH, _):\n      return True\n    case (lax.Precision.DEFAULT, _): # bfloat16\n      raise NotImplementedError(\"bfloat16 support not implemented for LSTM\")\n    case _:\n      raise ValueError(f\"Unexpected precision specifier value {precision}\")\n\n\n@partial(custom_vjp, nondiff_argnums=(5, 6, 7, 8, 9, 10))\ndef lstm(x: Array, h_0: Array, c_0: Array, weights: Array, seq_lengths: Array,\n         input_size: int, hidden_size: int, num_layers: int, dropout: float,\n         bidirectional: bool, precision: lax.PrecisionLike = None) -> tuple[Array, Array, Array]:\n  \"\"\"LSTM via CuDNN or HIPDNN (not-yet-supported).\n\n  Assume batch-first inputs.\n\n  Arguments:\n    x: (batch_size, max_seq_length, input_size)\n    h_0: (num_directions * num_layers, batch_size, hidden_size)\n    c_0: (num_directions * num_layers, batch_size, hidden_size)\n    weights: (num_params,) where num_params = get_num_params_in_lstm(...)\n    seq_lengths: (batch_size,)","sourceCodeStart":247,"sourceCodeEnd":283,"githubUrl":"https://github.com/jax-ml/jax/blob/1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb/jax/experimental/rnn.py#L247-L283","documentation":"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.","triggerScenarios":"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).","commonSituations":"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.","solutions":["Cast LSTM inputs/weights to float32 and pass precision=lax.Precision.HIGHEST to keep FP32 math","Use precision=lax.Precision.HIGH to allow TF32 (still FP32 storage) on Ampere+ GPUs","Fall back to a manual LSTM via lax.scan for bf16 training"],"exampleFix":"# before\ny, h_n, c_n = lstm(x_bf16, h0, c0, w_bf16, seq_lens, ..., precision=lax.Precision.DEFAULT)\n# after\ny, h_n, c_n = lstm(x_bf16.astype(jnp.float32), h0, c0, w.astype(jnp.float32), seq_lens, ..., precision=lax.Precision.HIGHEST)","handlingStrategy":"fallback","validationCode":"if precision is None or precision == lax.Precision.DEFAULT:\n    x, h0, c0, w = jax.tree.map(lambda a: a.astype(jnp.float32), (x, h0, c0, w))\n    precision = lax.Precision.HIGHEST","typeGuard":null,"tryCatchPattern":"try:\n    out = lstm(...)\nexcept NotImplementedError:\n    out = my_lax_scan_lstm(...)  # fp32 or manual fallback","preventionTips":["Keep cuDNN LSTM in float32/TF32","Centralize precision choice in one config"],"tags":["jax","lstm","rnn","bfloat16","cudnn","gpu"],"backgroundTag":"unsupported-dtype","analyzedSha":"1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb","analyzedAt":"2026-08-27T09:53:25.647Z","schemaVersion":2},"datasetVersion":"2026-08-27T13:17:12.746Z"}