{"record":{"id":"a1d41983fa3f913e","repo":"keras-team/keras","slug":"flash-attention-is-not-supported-in-your-current-j","errorCode":null,"errorMessage":"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","messagePattern":"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","errorType":"exception","errorClass":"ImportError","httpStatus":null,"severity":"error","filePath":"keras/src/backend/jax/nn.py","lineNumber":1472,"sourceCode":"    return psnr\n\n\ndef _can_use_flash_attention(query, key, value, bias, raise_error=False):\n    \"\"\"Verify the availability of flash attention.\"\"\"\n    try:\n        from jax._src.cudnn.fused_attention_stablehlo import _normalize_layout\n        from jax._src.cudnn.fused_attention_stablehlo import (\n            check_compute_capability,\n        )\n        from jax._src.cudnn.fused_attention_stablehlo import check_cudnn_version\n        from jax._src.cudnn.fused_attention_stablehlo import (\n            check_is_flash_attention,\n        )\n        from jax._src.cudnn.fused_attention_stablehlo import check_layout\n        from jax.nn import dot_product_attention as dot_product_attention\n    except ImportError:\n        if raise_error:\n            raise ImportError(\n                \"Flash attention is not supported in your current JAX version. \"\n                \"Please update it by following the official guide: \"\n                \"https://jax.readthedocs.io/en/latest/installation.html\"\n            )\n        return False\n\n    if jax.devices()[0].platform == \"tpu\":\n        return True\n    try:\n        # Check if cuDNN is installed and raise RuntimeError if cuDNN is not\n        # detected\n        cudnn_version = check_cudnn_version()\n        # Only support at least Ampere\n        if not check_compute_capability(\"8.0\"):\n            raise RuntimeError(\"Require at least Ampere arch to run\")\n\n        # Inspect inputs of `check_layout`\n        check_layout_params = list(","sourceCodeStart":1454,"sourceCodeEnd":1490,"githubUrl":"https://github.com/keras-team/keras/blob/7a34a03db60bf60042242d6a556fc3be119046a5/keras/src/backend/jax/nn.py#L1454-L1490","documentation":"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.","triggerScenarios":"Thrown at keras/src/backend/jax/nn.py:1472 when the library encounters an invalid state.","commonSituations":"See trigger scenarios.","solutions":[],"exampleFix":null,"handlingStrategy":null,"validationCode":null,"typeGuard":null,"tryCatchPattern":null,"preventionTips":[],"tags":[],"backgroundTag":null,"analyzedSha":"7a34a03db60bf60042242d6a556fc3be119046a5","analyzedAt":"2026-08-25T21:25:25.994Z","schemaVersion":2},"datasetVersion":"2026-08-26T02:17:13.382Z"}