class_docstr: Adds learned positional embeddings to the inputs.
  
  Attributes:
    posemb_init_space: Positional embedding initializer. Default value is taken
      from BERT.
    posemb_init_time: Positional embedding initializer. Default value is taken
      from BERT.
  
  Returns:
    Output with same shape as input.
class_docstr_tok: ["Adds","learned","positional","embeddings","to","the","inputs",".","Attributes",":","posemb_init_space",":","Positional","embedding","initializer",".","Default","value","is","taken","from","BERT",".","posemb_init_time",":","Positional","embedding","initializer",".","Default","value","is","taken","from","BERT",".","Returns",":","Output","with","same","shape","as","input","."]
code: def __call__(self, inputs: jnp.ndarray) -> jnp.ndarray:
      # Inputs.shape is [batch_size, time, space, hidden_dim].
      assert inputs.ndim == 4, (
          "Number of dimensions should be 4,"
          " but it is: %d" % inputs.ndim
      )
      _, time, space, hidden_dim = inputs.shape
      pos_emb_shape_space = (1, 1, space, hidden_dim)
      pos_emb_shape_time = (1, time, 1, hidden_dim)
      pe_spatial = self.param(
          "pos_embedding_space",
          self.posemb_init_space,
          pos_emb_shape_space,
          inputs.dtype,
      )
      pe_temporal = self.param(
          "pos_embedding_time",
          self.posemb_init_time,
          pos_emb_shape_time,
          inputs.dtype,
      )
      return inputs + pe_spatial + pe_temporal
  
file_path: scenic/scenic/projects/av_mae/model_utils.py
imports: from typing import Callable, Iterable, Optional, Sequence
  import flax.linen
  import jax
  import jax.numpy
  from scenic.model_lib.layers import attention_layers
  from scenic.projects.baselines import vit
label_desc: Use image node_id 1 for calling the AddFactorisedSpaceTimePositionEmbs obj's underlying member method code with example usage: obj.__call__(inputs) and returns: unknown