code: def test_cross_attention():
      torch.manual_seed(1234)
      for norm_first in [True, False]:
          m = StreamingTransformer(
              16,
              4,
              2,
              cross_attention=False,
              norm_first=norm_first,
              dropout=0.0,
              custom=True,
          )
          m_cross = StreamingTransformer(
              16,
              4,
              2,
              cross_attention=True,
              norm_first=norm_first,
              dropout=0.0,
              custom=True,
          )
          m_cross.load_state_dict(m.state_dict(), strict=False)
          x = torch.randn(2, 5, 16)
          cross_x = torch.randn(2, 3, 16)
          y_ref = m(x)
          y_cross_zero = m_cross(x, cross_attention_src=0 * cross_x)
          # With norm_first, the two should be exactly the same,
          # but with norm_first=False, we get 2 normalization in a row
          # and the epsilon value leads to a tiny change.
          atol = 0.0 if norm_first else 1e-6
          print((y_ref - y_cross_zero).norm() / y_ref.norm())
          assert torch.allclose(y_ref, y_cross_zero, atol=atol)
  
          # We now expect a difference even with a generous atol of 1e-2.
          y_cross = m_cross(x, cross_attention_src=cross_x)
          assert not torch.allclose(y_cross, y_cross_zero, atol=1e-2)
  
          with pytest.raises(AssertionError):
              _ = m_cross(x)
              _ = m(x, cross_attention_src=cross_x)
  
file_path: audiocraft/tests/modules/test_transformer.py
imports: from itertools import product
  import pytest
  import torch
  from audiocraft.modules.transformer import StreamingMultiheadAttention, StreamingTransformer, set_efficient_attention_backend
label_desc: Use image node_id 7 for calling a global function with example usage: test_cross_attention() without return types