code: def __call__(self, image: jnp.ndarray):
      x = image.astype(jnp.float32) / 255.0
      x = eg.Flatten()(x)
      x = eg.Linear(self.n1)(x)
      x = jax.nn.relu(x)
      x = eg.Linear(self.n2)(x)
      x = jax.nn.relu(x)
      x = eg.Linear(self.n1)(x)
      x = jax.nn.relu(x)
      x = eg.Linear(np.prod(image.shape[-2:]))(x)
      x = jax.nn.sigmoid(x) * 255
      x = x.reshape(image.shape)
  
      return x
  
file_path: elegy/examples/elegy/mnist_autoencoder.py
imports: import os
  from datetime import datetime
  from typing import Any, Generator, Mapping, Tuple
  import jax
  import jax.numpy
  import matplotlib.pyplot
  import numpy
  import optax
  import typer
  from datasets.load import load_dataset
  from tensorboardX.writer import SummaryWriter
  import elegy
label_desc: Use image node_id 2 for calling the MLP obj's underlying member method code with example usage: obj.__call__(image) and returns: x