code: def bernoulli_logpdf(logits, x):
      """Bernoulli log pdf of data x given logits."""
      return -jnp.sum(
          jnp.logaddexp(0.0, jnp.where(x, -1.0, 1.0) * logits)
      )
  
file_path: jax/examples/mnist_vae.py
imports: import os
  import time
  import matplotlib.pyplot
  import jax
  import jax.numpy
  from jax import jit, grad, lax, random
  from jax.example_libraries import optimizers
  from jax.example_libraries import stax
  from jax.example_libraries.stax import Dense, FanOut, Relu, Softplus
  from examples import datasets
label_desc: Use image node_id 3 for calling a global function with example usage: bernoulli_logpdf(logits, x) and returns: unknown