code: def forward(x):
      x = hk.Linear(64)(x)
      x = jax.nn.relu(x)
      x = hk.Linear(1)(x)
      return x
  
imports: import jax, optax
  import numpy
  import elegy
  import haiku
  import matplotlib.pyplot
label_desc: Use image node_id 1 for calling a global function with example usage: forward(x) and returns: x