keras-team/keras · error · ImportError

Flash attention is not supported in your current JAX version

Error message

Flash attention is not supported in your current JAX version. Please update it by following the official guide: https://jax.readthedocs.io/en/latest/installation.html

What it means

Error "Flash attention is not supported in your current JAX version. Please update it by following the official guide: https://jax.readthedocs.io/en/latest/installation.html" thrown in keras-team/keras.

Source

Thrown at keras/src/backend/jax/ops/nn.py:1473

    return psnr


def _can_use_flash_attention(query, key, value, bias, raise_error=False):
    """Verify the availability of flash attention."""
    try:
        from jax._src.cudnn.fused_attention_stablehlo import _normalize_layout
        from jax._src.cudnn.fused_attention_stablehlo import (
            check_compute_capability,
        )
        from jax._src.cudnn.fused_attention_stablehlo import check_cudnn_version
        from jax._src.cudnn.fused_attention_stablehlo import (
            check_is_flash_attention,
        )
        from jax._src.cudnn.fused_attention_stablehlo import check_layout
        from jax.nn import dot_product_attention as dot_product_attention
    except ImportError:
        if raise_error:
            raise ImportError(
                "Flash attention is not supported in your current JAX version. "
                "Please update it by following the official guide: "
                "https://jax.readthedocs.io/en/latest/installation.html"
            )
        return False

    if jax.devices()[0].platform == "tpu":
        return True
    try:
        # Check if cuDNN is installed and raise RuntimeError if cuDNN is not
        # detected
        cudnn_version = check_cudnn_version()
        # Only support at least Ampere
        if not check_compute_capability("8.0"):
            raise RuntimeError("Require at least Ampere arch to run")

        # Inspect inputs of `check_layout`
        check_layout_params = list(

View on GitHub (pinned to 7a34a03db6)

When it happens

Trigger: Thrown at keras/src/backend/jax/ops/nn.py:1473 when the library encounters an invalid state.

Common situations: See trigger scenarios.

AI-assisted analysis of keras-team/keras@7a34a03db6 (2026-08-25). Data as JSON: /api/errors/e72b544d4e71bcab. Report an issue: GitHub.