code: def binary_cross_entropy_with_mask(
      *, logits: jnp.ndarray, labels: jnp.ndarray, mask: jnp.ndarray
  ):
      """Binary cross entropy loss for unnormalized logits, with masked elements."""
      assert logits.shape == labels.shape == mask.shape
      assert len(logits.shape) == 2
  
      # To prevent propagation of NaNs during grad().
      # We mask over the loss for invalid targets later.
      labels = jnp.where(mask, labels, -1)
  
      # Numerically stable implementation of BCE loss.
      # This mimics TensorFlow's tf.nn.sigmoid_cross_entropy_with_logits().
      positive_logits = logits >= 0
      relu_logits = jnp.where(positive_logits, logits, 0)
      abs_logits = jnp.where(positive_logits, logits, -logits)
      return (
          relu_logits
          - (logits * labels)
          + (jnp.log(1 + jnp.exp(-abs_logits)))
      )
  
imports: import os
  from typing import Any, Dict, Iterable, Tuple, Optional
  from absl import logging
  from clu import checkpoint
  from clu import metric_writers
  from clu import metrics
  from clu import parameter_overview
  from clu import periodic_actions
  import flax
  import flax.core
  import flax.linen
  from flax.training import train_state
  import jax
  import jax.numpy
  import jraph
  import ml_collections
  import numpy
  import optax
  import sklearn.metrics
  import tensorflow
  import input_pipeline
  import models
label_desc: Use image node_id 3 for calling a global function with example usage: binary_cross_entropy_with_mask() and returns: unknown