class_docstr: Adds (optionally learned) positional embeddings to the inputs.
  
  Attributes:
    config: TransformerConfig dataclass containing hyperparameters.
    decode: whether to run in single-position autoregressive mode.
class_docstr_tok: ["Adds","(","optionally","learned",")","positional","embeddings","to","the","inputs",".","Attributes",":","config",":","TransformerConfig","dataclass","containing","hyperparameters",".","decode",":","whether","to","run","in","single-position","autoregressive","mode","."]
docstr: Applies AddPositionEmbs module.
  
  By default this layer uses a fixed sinusoidal embedding table. If a
  learned position embedding is desired, pass an initializer to
  posemb_init in the configuration.
  
  Args:
    inputs: input data.
    inputs_positions: input position indices for packed sequences.
  
  Returns:
    output: `(bs, timesteps, in_dim)`
docstr_tok: ["Applies","AddPositionEmbs","module",".","By","default","this","layer","uses","a","fixed","sinusoidal","embedding","table",".","If","a","learned","position","embedding","is","desired",",","pass","an","initializer","to","posemb_init","in","the","configuration",".","Args",":","inputs",":","input","data",".","inputs_positions",":","input","position","indices","for","packed","sequences",".","Returns",":","output",":","`","(","bs",",","timesteps",",","in_dim",")","`"]
code: def __call__(self, inputs, inputs_positions=None):
      """Applies AddPositionEmbs module.
  
      By default this layer uses a fixed sinusoidal embedding table. If a
      learned position embedding is desired, pass an initializer to
      posemb_init in the configuration.
  
      Args:
        inputs: input data.
        inputs_positions: input position indices for packed sequences.
  
      Returns:
        output: `(bs, timesteps, in_dim)`
      """
      config = self.config
      # inputs.shape is (batch_size, seq_len, emb_dim)
      assert inputs.ndim == 3, (
          "Number of dimensions should be 3,"
          " but it is: %d" % inputs.ndim
      )
      length = inputs.shape[1]
      pos_emb_shape = (1, config.max_len, inputs.shape[-1])
      if config.posemb_init is None:
          # Use a fixed (non-learned) sinusoidal position embedding.
          pos_embedding = sinusoidal_init(max_len=config.max_len)(
              None, pos_emb_shape, None
          )
      else:
          pos_embedding = self.param(
              "pos_embedding", config.posemb_init, pos_emb_shape
          )
      pe = pos_embedding[:, :length, :]
  
      # We use a cache position index for tracking decoding position.
      if self.decode:
          is_initialized = self.has_variable("cache", "cache_index")
          cache_index = self.variable(
              "cache",
              "cache_index",
              lambda: jnp.array(0, dtype=jnp.uint32),
          )
          if is_initialized:
              i = cache_index.value
              cache_index.value = i + 1
              _, _, df = pos_embedding.shape
              pe = lax.dynamic_slice(
                  pos_embedding, jnp.array((0, i, 0)), (1, 1, df)
              )
      if inputs_positions is None:
          # normal unpacked case:
          return inputs + pe
      else:
          # for packed data we need to use known position indices:
          return inputs + jnp.take(pe[0], inputs_positions, axis=0)
  
imports: from __future__ import annotations
  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 AddPositionEmbs obj's underlying member method code with example usage: obj.__call__(inputs, inputs_positions) and returns: unknown, unknown