code: def test_interleaved_multihead_self_attention(
      seq_len, batch_size, hidden, heads, side
  ):
      inputs = pt.rand((seq_len, batch_size, hidden))
  
      # test without attention masking
      mha = sockeye.layers.MultiHeadSelfAttention(
          hidden, heads, hidden, dropout=0.0
      )
      mha.train()
      assert not mha.kv_interleaved
      r_train, _ = mha(inputs, previous_states=None, mask=None)
      mha.eval()
      assert mha.kv_interleaved
      r_test, _ = mha(inputs, previous_states=None, mask=None)
      assert pt.allclose(r_train, r_test, atol=1e-06)
  
      # test with two types of attention masks (autoregressive, and valid_length based)
      if side == "decoder":
          # autoregressive mask. Shape: (len, len)
          mask = sockeye.transformer.AutoRegressiveMask()(
              inputs.transpose(0, 1)
          )
          mha.train()
          assert not mha.kv_interleaved
          r_train, _ = mha(inputs, previous_states=None, mask=mask)
          mha.eval()
          assert mha.kv_interleaved
          r_test, _ = mha(inputs, previous_states=None, mask=mask)
          assert pt.allclose(r_train, r_test, atol=1e-06)
      elif side == "encoder":
          all_source_length = pt.randint(1, seq_len + 1, (batch_size,))
          prepended_source_length = pt.full((batch_size,), 0)
          valid_length = pt.stack(
              (all_source_length, prepended_source_length), dim=1
          )
          # source attention mask. Shape: (batch * heads, 1, seq_len)
          mask = sockeye.layers.prepare_source_length_mask(
              valid_length, heads, seq_len
          )
          mask = mask.repeat(
              1, seq_len, 1
          )  # Shape: (batch * heads, seq_len, seq_len)
          mha.train()
          assert not mha.kv_interleaved
          r_train, _ = mha(inputs, previous_states=None, mask=mask)
          mha.eval()
          assert mha.kv_interleaved
          r_test, _ = mha(
              inputs, previous_states=None, mask=mask
          )  # Note: can also handle the mask repated on the qlen axis
          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 6 for calling a global function with example usage: test_interleaved_multihead_self_attention(seq_len, batch_size, hidden, heads, side) without return types