class_docstr: Tests for the EncoderDecoderModel.
class_docstr_tok: ["Tests","for","the","EncoderDecoderModel","."]
docstr: Tests metric_function by checking its output's format and validity.
docstr_tok: ["Tests","metric_function","by","checking","its","output","'s","format","and","validity","."]
code: def test_metric_function(self):
      """Tests metric_function by checking its output's format and validity."""
      model = FakeEncoderDecoderModel()
      batch, output = get_fake_batch_output()
      batch_replicated, outputs_replicated = (
          jax_utils.replicate(batch),
          jax_utils.replicate(output),
      )
  
      # Test metric function in the pmapped setup
      metrics_fn_pmapped = jax.pmap(
          model.get_metrics_fn(), axis_name="batch"
      )
      all_metrics = metrics_fn_pmapped(
          outputs_replicated, batch_replicated
      )
      # Check expected metrics exist in the output:
      expected_metrics_keys = ["accuracy", "loss", "perplexity"]
      self.assertSameElements(expected_metrics_keys, all_metrics.keys())
  
      # For each metric, check that it is a valid value.
      all_metrics = jax_utils.unreplicate(all_metrics)
      for k, v in all_metrics.items():
          self.is_valid(v[0], value_name=f"numerator of {k}")
          self.is_valid(v[1], value_name=f"denominator of {k}")
  
file_path: scenic/scenic/model_lib/base_models/tests/test_encoder_decoder_model.py
imports: from absl.testing import absltest
  from flax import jax_utils
  import jax
  import jax.numpy
  import ml_collections
  import numpy
  from scenic.model_lib.base_models import encoder_decoder_model
label_desc: Use image node_id 3 for calling the TestEncoderDecoderModel obj's underlying member method code with example usage: obj.test_metric_function() without return types