mirrors/rocm_jax
1
0
Fork 0
mirror of https://github.com/ROCm/jax.git synced 2025-04-19 05:16:06 +00:00
Code Issues Packages Projects Releases Wiki Activity

Fix lint failure in fused attention test.

Browse Source 
PiperOrigin-RevId: 606404377
This commit is contained in:
Peter Hawkins 2024-02-12 16:34:09 -08:00 committed by jax authors
parent 8e96c49f85
commit 7aadabdc03
1 changed files with 3 additions and 1 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

4
tests/fused_attention_stablehlo_test.py
View File

@ -14,7 +14,7 @@
from functools import partial
from absl.testing import absltest
from typing import Any, Optional
from typing import Optional
import os
os.environ['XLA_FLAGS'] = '--xla_gpu_enable_cudnn_fmha=true --xla_gpu_fused_attention_use_cudnn_rng=true'
@ -77,6 +77,8 @@ def sdpa_ref(query: Array,
if bias is not None:
attn_weights = attn_weights + bias.astype(attn_weights.dtype)
if mask is not None:
large_negative_number = jnp.asarray(
-0.7 * jnp.finfo(attn_weights.dtype).max, dtype=attn_weights.dtype)
attn_weights = jax.lax.select(mask, attn_weights, large_negative_number)
attn_weights = jax.nn.softmax(attn_weights)
if dropout_rate > 0.: