code: def test_attention_shape_inference(attention_cls):
      dummy_data = DummyData()
      attention = attention_cls(dummy_data.units, dummy_data.memory)
      attention_score = attention([dummy_data.query, dummy_data.state])
      assert len(attention_score) == 2
      assert attention_score[0].shape == (
          dummy_data.batch,
          dummy_data.timestep,
      )
      assert attention_score[1].shape == (
          dummy_data.batch,
          dummy_data.timestep,
      )
  
imports: import collections
  from packaging.version import Version
  import pytest
  import numpy
  import tensorflow
  from tensorflow_addons.seq2seq import attention_wrapper
  from tensorflow_addons.seq2seq import basic_decoder
  from tensorflow_addons.seq2seq import sampler
label_desc: Use image node_id 1 for calling a global function with example usage: test_attention_shape_inference(attention_cls) without return types