class_docstr: Test ObjectDetectionWithMatchingModel.
class_docstr_tok: ["Test","ObjectDetectionWithMatchingModel","."]
docstr: Test whether loss_function for the output and aux_outputs is same.
docstr_tok: ["Test","whether","loss_function","for","the","output","and","aux_outputs","is","same","."]
code: def test_auxiliary_loss_consistency(self):
      """Test whether loss_function for the output and aux_outputs is same."""
  
      model = MyObjectDetectionWithMatchingModel()
      outputs, batch, indices = fake_model_outputs_batch(num_boxes=4)
  
      # Test loss function in the pmapped setup:
      loss_function_pmapped = jax.pmap(
          model.loss_function, axis_name="batch"
      )
  
      indices_replicated = jax_utils.replicate(
          # Fake matching for the  final output +  2 aux outputs:
          [indices]
          * 3
      )
      outputs_replicated, batch_replicated = (
          jax_utils.replicate(outputs),
          jax_utils.replicate(batch),
      )
      _, metrics_dict = loss_function_pmapped(
          outputs_replicated, batch_replicated, indices_replicated
      )
  
      metrics_dict = jax_utils.unreplicate(metrics_dict)
  
      for key in ["loss_class", "loss_bbox", "loss_giou"]:
          for i in range(NUM_AUX_OUTPUTS):
              self.assertAlmostEqual(
                  metrics_dict[key + "_unscaled"],
                  metrics_dict[key + f"_aux_{i}_unscaled"],
                  places=5,
              )
              self.assertAlmostEqual(
                  metrics_dict[key],
                  metrics_dict[key + f"_aux_{i}"],
                  places=5,
              )
  
imports: from absl.testing import absltest
  from absl.testing import parameterized
  from flax import jax_utils
  import jax
  import jax.numpy
  import ml_collections
  import numpy
  from scenic.projects.baselines.detr import detr_base_model
label_desc: Use image node_id 6 for calling the TestObjectDetectionWithMatchingModel obj's underlying member method code with example usage: obj.test_auxiliary_loss_consistency() without return types