code: def __init__(self, cfg: Config, *, rngs: nnx.Rngs):
      sharding = cfg.sharding
      self.attn = Attention(cfg, rngs=rngs)
      self.mlp = MLP(cfg, rngs=rngs)
      self.scale1 = nnx.Param(
          jnp.ones((cfg.embed,), cfg.param_dtype), P(sharding.embed)
      )
      self.scale2 = nnx.Param(
          jnp.ones((cfg.embed,), cfg.param_dtype), P(sharding.embed)
      )
  
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 DecoderBlock object from inherited base classes: nnx with example: obj = DecoderBlock(cfg)