code: def get_net(x):
      def init(v):
          return dict(
              w_init=lambda *args: v * jnp.ones((1, 1)),
              b_init=lambda *args: v * 1.5 * jnp.ones((1,)),
          )
  
      h = basic.Linear(output_size=1, name="first_layer", **init(1.0))(
          x
      )
      h = basic.Linear(output_size=1, name="second_layer", **init(3.0))(
          h
      )
      return jnp.mean(h)
  
imports: import collections
  from collections.abc import Sequence
  import itertools
  import re
  import types
  from typing import Any, Callable
  from absl.testing import absltest
  from absl.testing import parameterized
  from haiku._src import basic
  from haiku._src import data_structures
  from haiku._src import filtering
  from haiku._src import transform
  import jax
  import jax.numpy
label_desc: Use image node_id 2 for calling a global function with example usage: get_net(x) and returns: jnp, dict