code: def __init__(self, cfg: Config, *, rngs: nnx.Rngs):
      sharding = cfg.sharding
  
      key = rngs.params()
      self.WQ = nnx.Param(
          dense_init(
              key,
              (cfg.embed, cfg.heads, cfg.depth),
              cfg.param_dtype,
              0,
              (1, 2),
          ),
          P(sharding.embed, sharding.heads, sharding.depth),
      )
      key = rngs.params()
      self.WK = nnx.Param(
          dense_init(
              key,
              (cfg.embed, cfg.heads, cfg.depth),
              cfg.param_dtype,
              0,
              (1, 2),
          ),
          P(sharding.embed, sharding.heads, sharding.depth),
      )
      key = rngs.params()
      self.WV = nnx.Param(
          dense_init(
              key,
              (cfg.embed, cfg.heads, cfg.depth),
              cfg.param_dtype,
              0,
              (1, 2),
          ),
          P(sharding.embed, sharding.heads, sharding.depth),
      )
      key = rngs.params()
      self.WO = nnx.Param(
          dense_init(
              key,
              (cfg.heads, cfg.depth, cfg.embed),
              cfg.param_dtype,
              (0, 1),
              2,
          ),
          P(sharding.heads, sharding.depth, sharding.embed),
      )
      # cache
      self.index = nnx.variable(
          "cache", jnp.array(0, dtype=jnp.int32), P()
      )
      self.key = nnx.variable(
          "cache",
          jnp.zeros(
              (cfg.batch, cfg.heads, cfg.depth, cfg.max_length),
              jnp.bfloat16,
          ),
          P(sharding.batch, sharding.heads, sharding.depth, None),
      )
      self = nnx.variable(
          "cache",
          jnp.zeros(
              (cfg.batch, cfg.heads, cfg.depth, cfg.max_length),
              jnp.bfloat16,
          ),
          P(sharding.batch, sharding.heads, sharding.depth, None),
      )
  
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 Attention object from inherited base classes: nnx with example: obj = Attention(cfg)