class_docstr: Encoder with facctorized self attention block.
  
  Attributes:
    mlp_dim: Dimension of the mlp on top of attention block.
    num_heads: Number of heads.
    temporal_dims: Number of temporal dimensions in the flattened input
    attention_kernel_initializer: Initializer to use for attention layers.
    dropout_rate: Dropout rate.
    attention_dropout_rate: Dropout for attention heads.
    droplayer_p: Probability of dropping a layer.
    attention_order: The order to do the attention. Choice of {time_space,
      space_time}.
    dtype: the dtype of the computation (default: float32).
class_docstr_tok: ["Encoder","with","facctorized","self","attention","block",".","Attributes",":","mlp_dim",":","Dimension","of","the","mlp","on","top","of","attention","block",".","num_heads",":","Number","of","heads",".","temporal_dims",":","Number","of","temporal","dimensions","in","the","flattened","input","attention_kernel_initializer",":","Initializer","to","use","for","attention","layers",".","dropout_rate",":","Dropout","rate",".","attention_dropout_rate",":","Dropout","for","attention","heads",".","droplayer_p",":","Probability","of","dropping","a","layer",".","attention_order",":","The","order","to","do","the","attention",".","Choice","of","{","time_space",",","space_time","}",".","dtype",":","the","dtype","of","the","computation","(","default",":","float32",")","."]
docstr: Applies Encoder1DBlock module.
code: def __call__(self, inputs: jnp.ndarray, *, deterministic: bool):
      """Applies Encoder1DBlock module."""
      b, thw, d = inputs.shape
      inputs = _reshape_to_time_space(inputs, self.temporal_dims)
      self_attention = functools.partial(
          nn.SelfAttention,
          num_heads=self.num_heads,
          kernel_init=self.attention_kernel_initializer,
          broadcast_dropout=False,
          dropout_rate=self.attention_dropout_rate,
          dtype=self.dtype,
      )
  
      if self.attention_order == "time_space":
          attention_axes = (1, 2)
      elif self.attention_order == "space_time":
          attention_axes = (2, 1)
      else:
          raise ValueError(
              f"Invalid attention order {self.attention_order}."
          )
  
      def _run_attention_on_axis(inputs, axis, two_d_shape):
          """Reshapes the input and run attention on the given axis."""
          inputs = model_utils.reshape_to_1d_factorized(
              inputs, axis=axis
          )
          x = nn.LayerNorm(
              dtype=self.dtype,
              name="LayerNorm_{}".format(_AXIS_TO_NAME[axis]),
          )(inputs)
          x = self_attention(
              name="MultiHeadDotProductAttention_{}".format(
                  _AXIS_TO_NAME[axis]
              )
          )(x, deterministic=deterministic)
          x = nn.Dropout(rate=self.dropout_rate)(x, deterministic)
          x = x + inputs
          return model_utils.reshape_to_2d_factorized(
              x, axis=axis, two_d_shape=two_d_shape
          )
  
      x = inputs
      two_d_shape = inputs.shape
      for axis in attention_axes:
          x = _run_attention_on_axis(x, axis, two_d_shape)
  
      # MLP block.
      x = jnp.reshape(x, [b, thw, d])
      y = nn.LayerNorm(dtype=self.dtype, name="LayerNorm_mlp")(x)
      y = attention_layers.MlpBlock(
          mlp_dim=self.mlp_dim,
          dtype=self.dtype,
          dropout_rate=self.dropout_rate,
          activation_fn=nn.gelu,
          kernel_init=nn.initializers.xavier_uniform(),
          bias_init=nn.initializers.normal(stddev=1e-6),
          name="MlpBlock",
      )(y, deterministic=deterministic)
      return x + y
  
imports: import functools
  from typing import Any, Optional, Callable, Sequence
  from absl import logging
  import flax.linen
  from flax.linen.linear import default_kernel_init
  from immutabledict import immutabledict
  import jax
  import jax.numpy
  import ml_collections
  import numpy
  from scenic.common_lib import video_utils
  from scenic.model_lib.base_models import base_model
  from scenic.model_lib.base_models import classification_model
  from scenic.model_lib.base_models import model_utils
  from scenic.model_lib.base_models.classification_model import ClassificationModel
  from scenic.model_lib.layers import attention_layers
  from scenic.model_lib.layers import nn_layers
  from scenic.projects.baselines import vit
  from scenic.projects.vivit import model_utils
label_desc: Use image node_id 1 for calling the EncoderFactorizedSelfAttentionBlock obj's underlying member method code with example usage: obj.__call__(inputs) and returns: unknown, model_utils