code: def __init__(self, cfg: Config, *, rngs: nnx.Rngs):
      sharding = cfg.sharding
      self.embed = nnx.Param(
          embed_init(
              rngs.params(),
              (cfg.vocab, cfg.embed),
              cfg.param_dtype,
              1,
              0,
          ),
          P(sharding.vocab, sharding.embed),
      )
      self.unembed = nnx.Param(
          dense_init(
              rngs.params(), (cfg.embed, cfg.vocab), jnp.float32, 0, 1
          ),
          P(sharding.embed, sharding.vocab),
      )
      self.scale1 = nnx.Param(
          jnp.ones((cfg.embed,), cfg.param_dtype), P(sharding.embed)
      )
  
      if cfg.scanned:
          self.layers = nnx.merge(
              jax.vmap(
                  lambda key: DecoderBlock(
                      cfg, rngs=nnx.Rngs(key)
                  ).split()
              )(jax.random.split(rngs.params(), cfg.layers))
          )
      else:
          self.layers = nnx.Sequence(
              DecoderBlock(cfg, rngs=rngs) for _ in range(cfg.layers)
          )
  
imports: import dataclasses
  import typing
  import jax
  import jax.numpy
  import numpy
  from jax.sharding import PartitionSpec
  from flax.experimental import nnx
label_desc: Use image node_id 1 to create a new Decoder object from inherited base classes: nnx with example: obj = Decoder(cfg)