def test_masking(self):
    """Check masked positions not contributing to input."""

    mention_mask = jnp.zeros(shape=(self.n_mentions))
    output, _ = self.model.init_with_output(
        self.rng,
        encoded_input=self.encoded_input,
        mention_batch_positions=self.mention_batch_positions,
        mention_start_positions=self.mention_start_positions,
        mention_end_positions=self.mention_end_positions,
        mention_mask=mention_mask,
        entity_embeddings=self.entity_embeddings,
    )

    self.assertTrue(jnp.allclose(output['encoded_output'], self.normed_input))
Check masked positions not contributing to input.
language/mentionmemory/modules/entity_attention_layer_test.py