jax

61374c92 - Fix error message in jax.nn.dot_product_attention when the inputs have different dtypes.