code: def __call__(self, x):
      x = nn.Dense(64)(x)
      x = jax.nn.relu(x)
      x = nn.Dense(1)(x)
      return x
  
file_path: elegy/examples/flax/toy_mlp.py
imports: import jax, optax
  import numpy
  import elegy
  import flax.linen
  import matplotlib.pyplot
label_desc: Use image node_id 1 for calling the MLP obj's underlying member method code with example usage: obj.__call__(x) and returns: x