def METHOD_NAME(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)