code: def train_step(state, batch, z_rng, latents):
      def loss_fn(params):
          recon_x, mean, logvar = models.model(latents).apply(
              {"params": params}, batch, z_rng
          )
  
          bce_loss = binary_cross_entropy_with_logits(
              recon_x, batch
          ).mean()
          kld_loss = kl_divergence(mean, logvar).mean()
          loss = bce_loss + kld_loss
          return loss
  
      grads = jax.grad(loss_fn)(state.params)
      return state.apply_gradients(grads=grads)
  
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 4 for calling a global function with example usage: train_step(state, batch, z_rng, latents) and returns: state, loss