def test_model_shape(
      self,
      text_length,
      n_mentions,
      n_linked_mentions,
      no_entity_attention,
  ):
    """Test model forward runs and produces expected shape."""

    config = copy.deepcopy(self.config)
    config['model_config']['encoder_config'][
        'no_entity_attention'] = no_entity_attention
    config = ml_collections.FrozenConfigDict(self.config)
    model_config = config.model_config
    encoder_config = model_config.encoder_config

    max_length = model_config.encoder_config.max_length
    preprocess_fn = eae_task.EaETask.make_preprocess_fn(config)
    collater_fn = eae_task.EaETask.make_collater_fn(config)

    raw_example = test_utils.gen_mention_pretraining_sample(
        text_length, n_mentions, n_linked_mentions, max_length=max_length)
    processed_example = preprocess_fn(raw_example)
    batch = {
        key: np.tile(value, (config.per_device_batch_size, 1))
        for key, value in processed_example.items()
    }
    batch = collater_fn(batch)
    batch = jax.tree.map(np.asarray, batch)

    model = eae_encoder.EaEEncoder(**model_config.encoder_config)
    init_rng = jax.random.PRNGKey(0)
    (encoded_output, loss_helpers, _), _ = model.init_with_output(
        init_rng,
        batch,
        deterministic=True,
        method=model.forward,
    )

    self.assertEqual(encoded_output.shape,
                     (config.per_device_batch_size, encoder_config.max_length,
                      encoder_config.hidden_size))

    self.assertEqual(loss_helpers['target_mention_encodings'].shape,
                     (config.max_mention_targets * config.per_device_batch_size,
                      encoder_config.entity_dim))
Test model forward runs and produces expected shape.
language/mentionmemory/encoders/eae_encoder_test.py