def test_model_shape(self, k_top):
    """Test batch memory attention layer output shape as expected."""

    model = batch_memory_attention_layer.BatchMemoryAttentionLayer(
        memory_key_dim=self.memory_key_dim,
        input_dim=self.input_dim,
        memory_update_type=self.memory_update_type,
        memory_update_config=self.memory_update_config,
        dtype=self.dtype,
        k_top=k_top,
        rows=self.rows,
    )

    rng = jax.random.PRNGKey(0)
    encoded_input = jnp.ones(
        shape=(self.bsz, self.seq_len, self.input_dim), dtype=self.dtype)

    mention_batch_positions = jnp.tile(
        jnp.arange(self.bsz).reshape(-1, 1), (1, 3)).reshape(-1)
    mention_start_positions = jnp.tile(jnp.asarray([0, 5, 10]), (self.bsz))
    mention_end_positions = jnp.tile(jnp.asarray([2, 7, 12]), (self.bsz))
    n_mentions = mention_start_positions.shape[-1]

    mention_mask = jnp.tile(jnp.asarray([1, 1, 1]), (self.bsz))

    memory_keys = jnp.ones((self.table_size, self.memory_key_dim),
                           dtype=self.dtype)
    memory_values = jnp.ones((self.table_size, self.memory_value_dim),
                             dtype=self.dtype)
    memory_entity_ids = jnp.arange(self.table_size)
    memory_mask = jnp.ones((self.table_size))

    (encoded_output, loss_helpers, _), _ = model.init_with_output(
        rng,
        encoding=encoded_input,
        mention_batch_positions=mention_batch_positions,
        mention_start_positions=mention_start_positions,
        mention_end_positions=mention_end_positions,
        mention_mask=mention_mask,
        memory_keys=memory_keys,
        memory_values=memory_values,
        memory_mask=memory_mask,
        memory_entity_ids=memory_entity_ids,
        deterministic=True,
    )

    attention_weights = loss_helpers['memory_attention_weights']
    entity_ids = loss_helpers['top_entity_ids']

    # Check shapes as expected
    self.assertSequenceEqual(encoded_output.shape,
                             (self.bsz, self.seq_len, self.input_dim))

    if k_top is None:
      return_shape = self.table_size
    else:
      return_shape = k_top

    self.assertSequenceEqual(attention_weights.shape,
                             (n_mentions, return_shape))

    self.assertSequenceEqual(entity_ids.shape, (n_mentions, return_shape))
Test batch memory attention layer output shape as expected.
language/mentionmemory/modules/batch_memory_attention_layer_test.py