code: def train_step(
      state: nnx.TrainState[VAE], x: jax.Array, key: jax.Array
  ):
      def loss_fn(params: nnx.State):
          rngs = nnx.Rngs(noise=jax.random.fold_in(key, state.step))
          logits, (updates, _) = state.apply(params)(x, rngs=rngs)
  
          losses = updates.extract(Loss)
          kl_loss = sum(jax.tree_util.tree_leaves(losses), 0.0)
          reconstruction_loss = jnp.mean(
              optax.sigmoid_binary_cross_entropy(logits, x)
          )
          # jax.debug.print("kl_loss={kl_loss}", kl_loss=kl_loss)
  
          loss = reconstruction_loss + 0.1 * kl_loss
          return loss
  
      loss, grads = jax.value_and_grad(loss_fn)(state.params)
      state = state.apply_gradients(grads=grads)
  
      return state, loss
  
imports: import typing
  from functools import partial
  import jax
  import jax.numpy
  import matplotlib.pyplot
  import numpy
  import optax
  from datasets import load_dataset
  from flax.experimental import nnx
label_desc: Use image node_id 1 for calling a global function with example usage: train_step(state, x, key) and returns: state,  loss, loss