class_docstr: Test various loss functions in model_utils.
class_docstr_tok: ["Test","various","loss","functions","in","model_utils","."]
docstr: Tests weighted_sigmoid_cross_entropy.
docstr_tok: ["Tests","weighted_sigmoid_cross_entropy","."]
code: def test_weighted_sigmoid_cross_entropy(self):
      """Tests weighted_sigmoid_cross_entropy."""
  
      logits = jnp.array([[1, 2, 3], [4, 5, 6]], dtype=jnp.float32)
      labels = jnp.array([[0, 1, 1], [1, 0, 1]], dtype=jnp.float32)
      sigmoid = jax.nn.sigmoid
      log = jnp.log
  
      loss = model_utils.weighted_sigmoid_cross_entropy(logits, labels)
      gt_loss = jnp.array(
          [
              [
                  -log(1 - sigmoid(1.0)),
                  -log(sigmoid(2.0)),
                  -log(sigmoid(3.0)),
              ],
              [
                  -log(sigmoid(4.0)),
                  -log(1 - sigmoid(5.0)),
                  -log(sigmoid(6.0)),
              ],
          ]
      ) / np.prod(labels.shape[:-1])
      self.assertSequenceAlmostEqual(
          loss.flatten(), gt_loss.sum().flatten(), places=3
      )
  
      example_weights = jnp.array([1.0, 0.0])
      loss = model_utils.weighted_sigmoid_cross_entropy(
          logits, labels, weights=example_weights
      )
      gt_loss = (
          jnp.array(
              [
                  [
                      -log(1 - sigmoid(1.0)),
                      -log(sigmoid(2.0)),
                      -log(sigmoid(3.0)),
                  ],
                  [0.0, 0.0, 0.0],
              ]
          )
          / example_weights.sum()
          + 1e-9
      )
      self.assertSequenceAlmostEqual(
          loss.flatten(), gt_loss.sum().flatten(), places=3
      )
  
      label_weights = jnp.array([1.0, 2.0, 3.0])
      loss = model_utils.weighted_sigmoid_cross_entropy(
          logits, labels, label_weights=label_weights
      )
      gt_loss = jnp.array(
          [
              [
                  -log(1 - sigmoid(1.0)),
                  -2 * log(sigmoid(2.0)),
                  -3 * log(sigmoid(3.0)),
              ],
              [
                  -log(sigmoid(4.0)),
                  -2 * log(1 - sigmoid(5.0)),
                  -3 * log(sigmoid(6.0)),
              ],
          ]
      ) / np.prod(labels.shape[:-1])
      self.assertSequenceAlmostEqual(
          loss.flatten(), gt_loss.sum().flatten(), places=3
      )
  
      loss = model_utils.weighted_sigmoid_cross_entropy(
          logits,
          labels,
          weights=example_weights,
          label_weights=label_weights,
      )
      gt_loss = (
          jnp.array(
              [
                  [
                      -log(1 - sigmoid(1.0)),
                      -2 * log(sigmoid(2.0)),
                      -3 * log(sigmoid(3.0)),
                  ],
                  [0.0, 0.0, 0.0],
              ]
          )
          / example_weights.sum()
          + 1e-9
      )
      self.assertSequenceAlmostEqual(
          loss.flatten(), gt_loss.sum().flatten(), places=3
      )
  
      # Label weights can actually be any shape that is broadcastable to the
      # shape of logits.
      label_weights = jnp.array([[1.0, 2.0, 3.0], [4.0, 5.0, 6.0]])
      loss = model_utils.weighted_sigmoid_cross_entropy(
          logits,
          labels,
          weights=example_weights,
          label_weights=label_weights,
      )
      gt_loss = (
          jnp.array(
              [
                  [
                      -log(1 - sigmoid(1.0)),
                      -2 * log(sigmoid(2.0)),
                      -3 * log(sigmoid(3.0)),
                  ],
                  [0.0, 0.0, 0.0],
              ]
          )
          / example_weights.sum()
          + 1e-9
      )
      self.assertSequenceAlmostEqual(
          loss.flatten(), gt_loss.sum().flatten(), places=3
      )
  
      with self.assertRaises(ValueError):
          label_weights = jnp.array([1.0, 2.0, 3.0, 4.0])
          loss = model_utils.weighted_sigmoid_cross_entropy(
              logits, labels, label_weights=label_weights
          )
  
imports: import itertools
  from absl.testing import absltest
  from absl.testing import parameterized
  from flax.training import common_utils
  import jax
  import jax.numpy
  import numpy
  from scenic.model_lib.base_models import model_utils
label_desc: Use image node_id 3 for calling the LossTest obj's underlying member method code with example usage: obj.test_weighted_sigmoid_cross_entropy() without return types