class_docstr: Adds (optionally learned) positional embeddings to the inputs.
  
  Args:
    config: TransformerConfig dataclass containing hyperparameters.
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: jax.Array, 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_embedding = self.pos_embedding
  
      if pos_embedding is None:
          # Use a fixed (non-learned) sinusoidal position embedding.
          pos_embedding = sinusoidal_init(max_len=config.max_len)(
              None, self.pos_emb_shape
          )
  
      # We use a cache position index for tracking decoding position.
      if flaglib.flags.get("decode", False):
          _, _, df = pos_embedding.shape
          # equivalent to pos_embedding[:, i:i+1] but traceable
          pos_embedding = lax.dynamic_slice(
              pos_embedding,
              jnp.array((0, self.cache_index, 0)),
              (1, 1, df),
          )
          self.cache_index += 1
      else:
          pos_embedding = pos_embedding[:, :length, :]
  
      if inputs_positions is None:
          # normal unpacked case:
          return inputs + pos_embedding
      else:
          # for packed data we need to use known position indices:
          return inputs + jnp.take(
              pos_embedding[0], inputs_positions, axis=0
          )
  
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 2 for calling the AddPositionEmbs obj's underlying member method code with example usage: obj.__call__(inputs, inputs_positions) and returns: unknown, unknown