class_docstr: Transformer Model for sequence to sequence translation.
  
  Attributes:
    config: TransformerConfig dataclass containing hyperparameters.
class_docstr_tok: ["Transformer","Model","for","sequence","to","sequence","translation",".","Attributes",":","config",":","TransformerConfig","dataclass","containing","hyperparameters","."]
code: def setup(self):
      config = self.config
  
      if config.share_embeddings:
          if config.output_vocab_size is not None:
              assert (
                  config.output_vocab_size == config.vocab_size
              ), "can't share embedding with different vocab sizes."
          self.shared_embedding = nn.Embed(
              num_embeddings=config.vocab_size,
              features=config.emb_dim,
              embedding_init=nn.initializers.normal(stddev=1.0),
          )
      else:
          self.shared_embedding = None
  
      self.encoder = Encoder(
          config=config, shared_embedding=self.shared_embedding
      )
      self.decoder = Decoder(
          config=config, shared_embedding=self.shared_embedding
      )
  
imports: from typing import Callable, Any, Optional
  from flax import linen
  from flax import struct
  from jax import lax
  import jax.numpy
  import numpy
label_desc: Use image node_id 1 for calling the Transformer obj's underlying member method code with example usage: obj.setup() without return types