def test_coref_resolution_loss_multiple_vs_single_devices(self):
    """Comparing coreference resolution loss on multiple vs single devices."""

    def compute_loss(mention_encodings, mention_batch_positions, mention_ids,
                     mention_target_is_masked):
      return mention_losses.coreference_resolution_loss(
          mention_encodings, mention_batch_positions, mention_ids,
          self.batch_size, 'dot', mention_target_is_masked, self.metrics_prefix)

    loss_sharded, metrics_sharded = jax.pmap(
        compute_loss, axis_name='batch')(self.mention_encodings_sharded,
                                         self.mention_batch_positions_sharded,
                                         self.mention_target_ids_sharded,
                                         self.mention_target_is_masked_sharded)

    mention_batch_positions_stacked = (
        self.mention_batch_positions_stacked.reshape(self.n_devices, -1))
    mention_batch_positions_stacked = mention_batch_positions_stacked.copy()
    mention_batch_positions_stacked += (
        np.expand_dims(np.arange(self.n_devices), 1) * self.batch_size)
    mention_batch_positions_stacked = mention_batch_positions_stacked.reshape(
        -1)
    loss_stacked, metrics_stacked = compute_loss(
        self.mention_encodings_stacked, mention_batch_positions_stacked,
        self.mention_target_ids_stacked, self.mention_target_is_masked_stacked)

    loss_sharded = jnp.sum(loss_sharded)
    metrics_sharded = jax.tree.map(jnp.sum, metrics_sharded)

    self.assertAlmostEqual(loss_sharded, loss_stacked, places=2)
    for metric_group_name in metrics_stacked:
      for metric_name in metrics_stacked[metric_group_name]:
        self.assertAlmostEqual(
            metrics_sharded[metric_group_name][metric_name],
            metrics_stacked[metric_group_name][metric_name],
            places=2)
Comparing coreference resolution loss on multiple vs single devices.
language/mentionmemory/modules/mention_losses_test.py