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 __init__(self, dim: int, *, n_layers: int, rngs: nnx.Rngs):
      self.n_layers = n_layers
      # fork Rngs, split keys into `n_layers`
      keys = rngs.fork(n_layers)
  
      def create_block(keys):
          # create Block instance and return its split
          return Block(dim, rngs=nnx.Rngs(keys)).split()
  
      # call vmap over create_block, passing the split `params` key
      # and immediately merge to get a Block instance
      self.layers = nnx.merge(jax.vmap(create_block)(keys))
  
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 1 to create a new ScanMLP object from inherited base classes: nnx with example: obj = ScanMLP(dim)