class_docstr: Tests for the a fake regression model.
class_docstr_tok: ["Tests","for","the","a","fake","regression","model","."]
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 = FakeRegressionModel()
      batch, predictions = get_fake_batch_and_predictions()
      batch_replicated, predictions_replicated = (
          jax_utils.replicate(batch),
          jax_utils.replicate(predictions),
      )
  
      metrics_fn_pmapped = jax.pmap(
          model.get_metrics_fn(), axis_name="batch"
      )
      all_metrics = metrics_fn_pmapped(
          predictions_replicated, batch_replicated
      )
      expected_metrics_keys = ["mean_squared_error"]
      self.assertSameElements(expected_metrics_keys, all_metrics.keys())
  
      all_metrics = jax_utils.unreplicate(all_metrics)
      self.assertLen(all_metrics, 1)
  
      mse_sum_count = all_metrics["mean_squared_error"]
      # (|[0, 1, 0, 0]|^2 + |[0, 0, 0, 0|^2 + |[1, 3, 0, 0]|^2) = 11
      self.assertAlmostEqual(mse_sum_count[0], 11.0)
      self.assertEqual(mse_sum_count[1], 3)
  
file_path: scenic/scenic/model_lib/base_models/tests/test_regression_model.py
imports: from absl.testing import absltest
  from flax import jax_utils
  import jax
  import jax.numpy
  import ml_collections
  from scenic.model_lib.base_models import regression_model
label_desc: Use image node_id 2 for calling the TestRegressionModel obj's underlying member method code with example usage: obj.test_metric_function() without return types