code: def forward(
      state: nnx.TrainState[VAE], x: jax.Array, key: jax.Array
  ) -> jax.Array:
      rngs = nnx.Rngs(noise=key)
      y_pred = state.apply("params")(x, rngs=rngs)[0]
      return jax.nn.sigmoid(y_pred)
  
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 2 for calling a global function with example usage: forward(state, x, key) and returns: jax