class_docstr: Tests for the MultiLabelClassificationModel.
class_docstr_tok: ["Tests","for","the","MultiLabelClassificationModel","."]
docstr: Tests a masked loss_function by comparing different canonical masks.
docstr_tok: ["Tests","a","masked","loss_function","by","comparing","different","canonical","masks","."]
code: def test_loss_function_masked(self):
      """Tests a masked loss_function by comparing different canonical masks."""
      array_size = (BATCH_SIZE, 50, NUM_CLASSES)
  
      model = FakeMultiLabelClassificationModel()
      batch, output = get_fake_batch_output(array_size=array_size)
  
      # Unmasked loss
      loss_value_unmasked = model.loss_function(output, batch)
  
      # Mask with only ones (so effectively no mask).
      batch["batch_mask"] = jnp.ones((BATCH_SIZE, 50))
      loss_value_masked = model.loss_function(output, batch)
  
      self.assertAlmostEqual(
          float(loss_value_unmasked), float(loss_value_masked)
      )
  
      # Extend the batch with random outputs and labels, but mask them with 0's.
      batch_extended = {
          "label": jnp.concatenate(
              (batch["label"], np.random.randint(2, size=array_size)),
              axis=1,
          ),
          "batch_mask": jnp.concatenate(
              (batch["batch_mask"], np.zeros((BATCH_SIZE, 50))), axis=1
          ),
      }
      output_extended = jnp.concatenate(
          (output, np.random.random(size=array_size)), axis=1
      )
      loss_value_extended = model.loss_function(
          output_extended, batch_extended
      )
  
      # Test with `places=3` due to JAX issue: github.com/google/jax/issues/6553
      # TODO(robromijnders): follow up with JAX issue and remove `places=3`.
      self.assertAlmostEqual(
          float(loss_value_masked), float(loss_value_extended), places=3
      )
  
file_path: scenic/scenic/model_lib/base_models/tests/test_multilabel_classification_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 multilabel_classification_model
label_desc: Use image node_id 3 for calling the TestMultiLabelClassificationModel obj's underlying member method code with example usage: obj.test_loss_function_masked() without return types