code: def test_interleaved_multihead_attention(
      qlen, kvlen, batch_size, hidden, heads
  ):
      queries_pt = pt.rand((qlen, batch_size, hidden))
      memory_pt = pt.rand((kvlen, batch_size, hidden))
  
      # test without mask
      mha = sockeye.layers.MultiHeadAttention(
          hidden, heads, hidden, dropout=0.0, depth_key_value=hidden
      )
      mha.train()
      assert not mha.kv_interleaved
      r_train = mha(
          queries_pt, memory_pt, mask=None, projected_memory_kv=None
      )
      mha.eval()
      assert mha.kv_interleaved
      r_test = mha(
          queries_pt, memory_pt, mask=None, projected_memory_kv=None
      )
      assert pt.allclose(r_train, r_test, atol=1e-06)
  
      # test with mask
      all_source_length = pt.randint(1, kvlen + 1, (batch_size,))
      prepended_source_length = pt.full((batch_size,), 0)
      valid_length = pt.stack(
          (all_source_length, prepended_source_length), dim=1
      )
      mask = sockeye.layers.prepare_source_length_mask(
          valid_length, heads, kvlen
      )
      mask = mask.repeat(
          1, qlen, 1
      )  # Shape: (batch *h heads, qlen, kvlen)
      mha.train()
      assert not mha.kv_interleaved
      r_train = mha(
          queries_pt, memory_pt, mask=mask, projected_memory_kv=None
      )
      mha.eval()
      assert mha.kv_interleaved
      r_test = mha(
          queries_pt, memory_pt, mask=mask, projected_memory_kv=None
      )
      assert pt.allclose(r_train, r_test, atol=1e-06)
  
file_path: sockeye/test/unit/test_layers.py
imports: from math import pow, sqrt
  import numpy
  import pytest
  import torch
  import sockeye.constants
  import sockeye.layers
  import sockeye.transformer
  from sockeye.knn import KNNConfig, FaissIndexBuilder
label_desc: Use image node_id 5 for calling a global function with example usage: test_interleaved_multihead_attention(qlen, kvlen, batch_size, hidden, heads) without return types