code: def binary_cross_entropy_with_logits(logits, labels):
      logits = nn.log_sigmoid(logits)
      return -jnp.sum(
          labels * logits + (1.0 - labels) * jnp.log(-jnp.expm1(logits))
      )
  
file_path: flax/examples/vae/train.py
imports: from absl import logging
  from flax import linen
  import input_pipeline
  import models
  import utils
  from flax.training import train_state
  import jax
  from jax import random
  import jax.numpy
  import ml_collections
  import optax
  import tensorflow_datasets
label_desc: Use image node_id 2 for calling a global function with example usage: binary_cross_entropy_with_logits(logits, labels) and returns: unknown