code: def mlp_vmap(
      scope: Scope,
      x: Array,
      sizes: Sequence[int] = (8, 1),
      act_fn: Callable[[Array], Array] = nn.relu,
      share_params: bool = False,
  ):
      if share_params:
          dense_vmap = lift.vmap(
              nn.dense,
              in_axes=(0, None),
              variable_axes={"params": None},
              split_rngs={"params": False},
          )
      else:
          dense_vmap = lift.vmap(
              nn.dense,
              in_axes=(0, None),
              variable_axes={"params": 0},
              split_rngs={"params": True},
          )
  
      # hidden layers
      for size in sizes[:-1]:
          x = scope.child(dense_vmap, prefix="hidden_")(x, size)
          x = act_fn(x)
  
      # output layer
      return scope.child(dense_vmap, "out")(x, sizes[-1])
  
file_path: flax/tests/core/design/core_vmap_test.py
imports: from typing import Callable, Sequence
  import jax
  from absl.testing import absltest
  from jax import numpy
  from jax import random
  from flax.core import Array, Scope, init, lift, nn, unfreeze
label_desc: Use image node_id 1 for calling a global function with example usage: mlp_vmap(scope, x, sizes, act_fn, share_params) and returns: scope