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

CI: fix mypy error

Browse Source
This commit is contained in:
Jake VanderPlas 2024-05-20 13:23:15 -07:00
parent 06d2e489eb
commit 329ab036ee
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

@ -986,7 +986,7 @@ def dot_product_attention(query: Array,
has_bias = bias is not None
has_mask = mask is not None
has_dbias = has_bias and is_training and \
should_export_dbias(bias.shape, query.shape, layout)
should_export_dbias(bias.shape, query.shape, layout) # type: ignore[union-attr]
variadic_args = (has_bias, has_mask, has_dbias)
if bias is None:
bias = jnp.zeros(0, dtype=query.dtype)