class_docstr: An MLP that uses `vmap` during `__init__` to create a Block instance
  with an additional `layer` axis, and `scan` during `__call__` to apply
  the sequence of layers iteratively over the input / output `x`.
class_docstr_tok: ["An","MLP","that","uses","`","vmap","`","during","`","__init__","`","to","create","a","Block","instance","with","an","additional","`","layer","`","axis",",","and","`","scan","`","during","`","__call__","`","to","apply","the","sequence","of","layers","iteratively","over","the","input","\/","output","`","x","`","."]
code: def __call__(self, x: jax.Array, *, rngs: nnx.Rngs) -> jax.Array:
      # fork Rngs, split keys into `n_layers`
      keys = rngs.fork(self.n_layers)
      # split Module to get params
      params, static = self.layers.split(nnx.Param)
  
      def scan_fn(
          x: jax.Array,
          inputs: Tuple[nnx.State, dict[str, nnx.RngStream]],
      ) -> Tuple[jax.Array, nnx.State]:
          params, keys = inputs
          # merge back Module and Rngs
          module = static.merge(params)
          # forward pass
          x = module(x, rngs=nnx.Rngs(keys))
          # split state and return
          params, _ = module.split(nnx.Param)
          return x, params
  
      # call scan passing x as the carry, and params + keys as the input
      x, params = jax.lax.scan(scan_fn, x, (params, keys))
      # update layers state and return
      self.layers.update(params)
      return x
  
file_path: flax/flax/experimental/nnx/examples/toy_examples/06_scan_over_layers.py
imports: from typing import Tuple
  import jax
  import jax.numpy
  from flax.experimental import nnx
label_desc: Use image node_id 2 for calling the ScanMLP obj's underlying member method code with example usage: obj.__call__(x) and returns: x, x,  params