class_docstr: Tests for the a fake regression model.
class_docstr_tok: ["Tests","for","the","a","fake","regression","model","."]
docstr: Tests loss_function by checking its output's validity.
docstr_tok: ["Tests","loss_function","by","checking","its","output","'s","validity","."]
code: def test_loss_function(self):
      """Tests loss_function by checking its output's validity."""
      model = FakeRegressionModel()
      batch, predictions = get_fake_batch_and_predictions()
      batch_replicated, predictions_replicated = (
          jax_utils.replicate(batch),
          jax_utils.replicate(predictions),
      )
  
      # Test loss function in the pmapped setup:
      loss_function_pmapped = jax.pmap(
          model.loss_function, axis_name="batch"
      )
      total_loss = loss_function_pmapped(
          predictions_replicated, batch_replicated
      )
      total_loss = jax_utils.unreplicate(total_loss)
      # Loss =  1/3 * (|[0, 1, 0, 0]|^2 + |[0, 0, 0, 0|^2 + |[1, 3, 0, 0]|^2)
      self.assertAlmostEqual(total_loss, 11 / 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 1 for calling the TestRegressionModel obj's underlying member method code with example usage: obj.test_loss_function() without return types