bidirectional_cross_attention_jax
