def test_loss_fn(self, n_mentions, apply_mlp=False):
    """Test loss function runs and produces expected values."""
    config = copy.deepcopy(self.config)
    config['model_config']['apply_mlp'] = apply_mlp
    raw_batch = self._gen_raw_batch(n_mentions)
    batch = self.collater_fn(raw_batch)
    batch = jax.tree.map(np.asarray, batch)

    loss_fn = text_classifier.TextClassifier.make_loss_fn(config)
    _, metrics, auxiliary_output = loss_fn(
        model_config=self.model_config,
        model_params=self.init_parameters['params'],
        model_vars={},
        batch=batch,
        deterministic=True)

    self.assertEqual(metrics['agg']['denominator'],
                     config.per_device_batch_size)
    features = self.postprocess_fn(batch, auxiliary_output)
    # Check features are JSON-serializable
    json.dumps(features)
    # Check features match the original batch
    for key in batch.keys():
      self.assertArrayEqual(np.array(features[key]), batch[key])
Test loss function runs and produces expected values.
language/mentionmemory/tasks/text_classifier_test.py