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 mypy error

Browse Source
This commit is contained in:
Jake VanderPlas 2024-07-07 19:45:50 -07:00
parent ceea8dc3fe
commit f851099649
1 changed files with 1 additions and 1 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

2
jax/_src/cudnn/fused_attention_stablehlo.py
View File

@ -980,7 +980,7 @@ def dot_product_attention(query: Array,
large_negative_number = get_large_negative_number(query.dtype)
mask = jnp.where(mask, jnp.asarray(0, query.dtype), large_negative_number)
# reshape mask to have 4D shape
mask = mask.reshape((1,) * (4 - len(mask.shape)) + mask.shape)
mask = mask.reshape((1,) * (4 - len(mask.shape)) + mask.shape) # type: ignore[union-attr]
# combine bias and mask
if bias is None: