code: def mlp_scan(scope: Scope, xs: Array, share_params: bool = False):
      scope.variable("counter", "i", jnp.zeros, ())
  
      def body_fn(scope, c, x):
          counter = scope.variable("counter", "i", jnp.zeros, ())
          counter.value += 1
          x = scope.child(nn.dense)(x, 1)
          return c, x
  
      if share_params:
          _, ys = lift.scan(
              body_fn,
              variable_carry="counter",
              variable_broadcast="params",
              split_rngs={"params": False},
          )(scope, (), xs)
      else:
          _, ys = lift.scan(
              body_fn,
              variable_carry="counter",
              variable_axes={"params": 0},
              split_rngs={"params": True},
          )(scope, (), xs)
  
      # output layer
      return ys
  
file_path: flax/tests/core/design/core_scan_test.py
imports: 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_scan(scope, xs, share_params) and returns: ys, c,  x