def sigmoid_cross_entropy_with_logits(*, labels: jnp.ndarray,
                                      logits: jnp.ndarray) -> jnp.ndarray:
  """Sigmoid cross entropy loss.

  We use a stable formulation that is equivalent to the one used in TensorFlow.
  The following derivation shows how we arrive at the formulation:

  .. math::
        z * -log(sigmoid(x)) + (1 - z) * -log(1 - sigmoid(x))
      = z * -log(1 / (1 + exp(-x))) + (1 - z) * -log(exp(-x) / (1 + exp(-x)))
      = z * log(1 + exp(-x)) + (1 - z) * (-log(exp(-x)) + log(1 + exp(-x)))
      = z * log(1 + exp(-x)) + (1 - z) * (x + log(1 + exp(-x))
      = (1 - z) * x + log(1 + exp(-x))
      = x - x * z + log(1 + exp(-x))

  For x < 0, the following formula is more stable:
  .. math::
        x - x * z + log(1 + exp(-x))
      = log(exp(x)) - x * z + log(1 + exp(-x))
      = - x * z + log(1 + exp(x))

  We combine the two cases (x<0, x>=0) into one formula as follows:
  .. math::
      max(x, 0) - x * z + log(1 + exp(-abs(x)))

  This function is vmapped, so it is written for a single example, but can
  handle a batch of examples.

  Args:
    labels: The correct labels.
    logits: The output logits.

  Returns:
    The binary cross entropy loss for each given logit.
  """
  # The final formulation is: max(x, 0) - x * z + log(1 + exp(-abs(x)))
  # To allow computing gradients at zero, we define custom versions of max and
  # abs functions just like in tf.nn.sigmoid_cross_entropy_with_logits.
  zeros = jnp.zeros_like(logits, dtype=logits.dtype)
  condition = (logits >= zeros)
  relu_logits = jnp.where(condition, logits, zeros)
  neg_abs_logits = jnp.where(condition, -logits, logits)
  return relu_logits - logits * labels + jnp.log1p(jnp.exp(neg_abs_logits))
Sigmoid cross entropy loss.

  We use a stable formulation that is equivalent to the one used in TensorFlow.
  The following derivation shows how we arrive at the formulation:

  .. math::
        z * -log(sigmoid(x)) + (1 - z) * -log(1 - sigmoid(x))
      = z * -log(1 / (1 + exp(-x))) + (1 - z) * -log(exp(-x) / (1 + exp(-x)))
      = z * log(1 + exp(-x)) + (1 - z) * (-log(exp(-x)) + log(1 + exp(-x)))
      = z * log(1 + exp(-x)) + (1 - z) * (x + log(1 + exp(-x))
      = (1 - z) * x + log(1 + exp(-x))
      = x - x * z + log(1 + exp(-x))

  For x < 0, the following formula is more stable:
  .. math::
        x - x * z + log(1 + exp(-x))
      = log(exp(x)) - x * z + log(1 + exp(-x))
      = - x * z + log(1 + exp(x))

  We combine the two cases (x<0, x>=0) into one formula as follows:
  .. math::
      max(x, 0) - x * z + log(1 + exp(-abs(x)))

  This function is vmapped, so it is written for a single example, but can
  handle a batch of examples.

  Args:
    labels: The correct labels.
    logits: The output logits.

  Returns:
    The binary cross entropy loss for each given logit.
  
videogvt/train_lib/losses.py