code: def __init__(
      self,
      embed_dim,
      num_heads,
      ff_dim,
      prenorm=False,
      act=lambda x: x.relu(),
      dropout=0.1,
  ):
      assert (
          embed_dim % num_heads == 0
      ), "embed_dim must be divisible by num_heads"
  
      self.num_heads = num_heads
      self.head_size = embed_dim // num_heads
      self.prenorm, self.act = prenorm, act
      self.dropout = dropout
  
      self.query = (
          Tensor.scaled_uniform(embed_dim, embed_dim),
          Tensor.zeros(embed_dim),
      )
      self.key = (
          Tensor.scaled_uniform(embed_dim, embed_dim),
          Tensor.zeros(embed_dim),
      )
      self.value = (
          Tensor.scaled_uniform(embed_dim, embed_dim),
          Tensor.zeros(embed_dim),
      )
  
      self.out = (
          Tensor.scaled_uniform(embed_dim, embed_dim),
          Tensor.zeros(embed_dim),
      )
  
      self.ff1 = (
          Tensor.scaled_uniform(embed_dim, ff_dim),
          Tensor.zeros(ff_dim),
      )
      self.ff2 = (
          Tensor.scaled_uniform(ff_dim, embed_dim),
          Tensor.zeros(embed_dim),
      )
  
      self.ln1 = (Tensor.ones(embed_dim), Tensor.zeros(embed_dim))
      self.ln2 = (Tensor.ones(embed_dim), Tensor.zeros(embed_dim))
  
imports: from tinygrad import Tensor
label_desc: Use image node_id 1 to create a new TransformerBlock object  with example: obj = TransformerBlock(embed_dim, num_heads, ff_dim, prenorm, act, dropout)