code: def enc_dec_mix(enc, dec, enc_mask, dec_mask, n_head=16):
      # enc = batch, seq, feats
      # dec = batch, seq, feats
      with tf.compat.v1.variable_scope("enc_dec_attn"):
          batch, dec_seq, feats = shape_list(dec)
          enc_seq = shape_list(enc)[1]
  
          enc_proj = normal_1d_conv_block(
              enc,
              1,
              "enc_proj",
              use_fp16=enc.dtype == tf.float16,
              output_dim=feats * 2,
          )
          dec_proj = normal_1d_conv_block(
              dec,
              1,
              "dec_proj",
              use_fp16=enc.dtype == tf.float16,
              output_dim=feats,
          )
          k, v = tf.split(enc_proj, 2, 2)
          q = dec_proj
          q = split_heads(q, n_head)
          k = split_heads(k, n_head, k=True)
          v = split_heads(v, n_head)
          w = tf.matmul(q, k)
          w = w * tf.math.rsqrt(cast_maybe(dec_seq, tf.float32))
  
          enc_mask = tf.reshape(
              tf.sequence_mask(
                  enc_mask, maxlen=enc_seq, dtype=enc.dtype
              ),
              [batch, 1, 1, enc_seq],
          )
          dec_mask = tf.reshape(
              tf.sequence_mask(
                  dec_mask, maxlen=dec_seq, dtype=enc.dtype
              ),
              [batch, 1, dec_seq, 1],
          )
          m = enc_mask * dec_mask
          w = w * m + -1e9 * (1 - m)
          w = tf.nn.softmax(w)
  
      return merge_heads(tf.matmul(w, v))
  
file_path: finetune/finetune/base_models/oscar/featurizer.py
imports: import tensorflow
  from finetune.base_models.gpt.featurizer import dropout, embed, split_heads, merge_heads
  from finetune.util.shapes import shape_list, lengths_from_eos_idx
  from finetune.custom_ops import recursive_agg
  from finetune.base_models.gpt.featurizer import get_pos_values
  from finetune.optimizers.recompute_grads import recompute_grad
  import functools
label_desc: Use image node_id 9 for calling a global function with example usage: enc_dec_mix(enc, dec, enc_mask, dec_mask, n_head) and returns: merge_heads