使用self.param方法和posemb_init函数初始化位置

class AddPositionEmbs(nn.Module):
  """Adds learned positional embeddings to the inputs.

  Attributes:
    posemb_init: positional embedding initializer.
  """

  posemb_init: Callable[[PRNGKey, Shape, Dtype], Array]
  #PRNGKey：一个用于生成随机数的伪随机数生成器的密钥。
  #Shape：一个元组，表示要初始化的位置嵌入的形状。
  #Dtype：位置嵌入的数据类型。
  #你是在声明posemb_init是一个类型为Callable[[PRNGKey, Shape, Dtype], Array]的变量。
  #这意味着posemb_init应该是一个可调用对象（比如函数），它接受三个参数（一个PRNGKey、一个Shape和一个Dtype），并返回一个Array。
  param_dtype: Dtype = jnp.float32
  #这段代码为 param_dtype 定义了一个默认值 jnp.float32，它将被用作参数的数据类型（dtype）。

  @nn.compact
  def __call__(self, inputs):
    """Applies the AddPositionEmbs module.

    Args:
      inputs: Inputs to the layer.

    Returns:
      Output tensor with shape `(bs, timesteps, in_dim)`.
      bs 表示批量大小（batch size），timesteps 表示时间步（或序列长度），in_dim 表示输入维度。
    """
    # inputs.shape is (batch_size, seq_len, emb_dim).
    #bz:表示每个批次（batch）中图像的数量。
    #seq_len:seq_len 表示了输入序列的长度，即图像被划分成了多少个序列元素(patch)。
    #emb_dim:图片展开的维度，例如16*16*3
    pos_emb_shape = (1, inputs.shape[1], inputs.shape[2])
    #pos_emb_shape这个形状通常用于初始化位置嵌入的张量，以确保位置嵌入与输入张量的维度相匹配。
    pe = self.param(
        'pos_embedding', self.posemb_init, pos_emb_shape, self.param_dtype)
    return inputs + pe