class_docstr: Transformer Model Decoder for sequence to sequence translation.
  
  Args:
    config: TransformerConfig dataclass containing hyperparameters.
    shared_embedding: a shared embedding layer to use.
code: def __init__(
      self,
      config: TransformerConfig,
      shared_embedding: nnx.Embed | None = None,
      *,
      rngs: nnx.Rngs,
  ):
      self.config = config
      self.shared_embedding = shared_embedding
  
      # Target Embedding
      if self.shared_embedding is None:
          self.output_embed = nnx.Embed(
              num_embeddings=config.output_vocab_size,
              features=config.emb_dim,
              embedding_init=nnx.with_partitioning(
                  nnx.initializers.normal(stddev=1.0),
                  config.axis_rules("vocab", "embed"),
              ),
              rngs=rngs,
          )
      else:
          self.output_embed = self.shared_embedding
  
      self.posembed_output = AddPositionEmbs(config=config, rngs=rngs)
      self.dropout = nnx.Dropout(rate=config.dropout_rate)
      for idx in range(config.num_layers):
          layer = EncoderDecoder1DBlock(config=config, rngs=rngs)
          setattr(self, f"encoderdecoderblock_{idx}", layer)
  
      self.encoderdecoder_norm = nnx.LayerNorm(
          num_features=config.emb_dim,
          dtype=config.dtype,
          bias_init=nnx.with_partitioning(
              nnx.initializers.zeros_init(), config.axis_rules("embed")
          ),
          scale_init=nnx.with_partitioning(
              nnx.initializers.ones_init(), config.axis_rules("embed")
          ),
          rngs=rngs,
      )
      if not config.logits_via_embedding:
          self.logitdense = nnx.Linear(
              in_features=config.emb_dim,
              out_features=config.output_vocab_size,
              dtype=config.dtype,
              kernel_init=nnx.with_partitioning(
                  config.kernel_init,
                  config.axis_rules("embed", "vocab"),
              ),
              bias_init=nnx.with_partitioning(
                  config.bias_init, config.axis_rules("vocab")
              ),
              rngs=rngs,
          )
      else:
          self.logitdense = None
  
imports: from __future__ import annotations
  import dataclasses
  from typing import Any, Optional
  import jax
  import jax.numpy
  import numpy
  from jax import lax
  from flax.experimental import nnx
  from flax.experimental.nnx.nnx import flaglib
  from flax.experimental.nnx.examples.lm1b.configs import default
label_desc: Use image node_id 1 to create a new Decoder object from inherited base classes: nnx with example: obj = Decoder(config, shared_embedding)