class_docstr: ViT model for Video with factorized space-time attention.
class_docstr_tok: ["ViT","model","for","Video","with","factorized","space-time","attention","."]
code: def __call__(
      self, x: jnp.ndarray, *, train: bool, debug: bool = False
  ):
      del debug
      x, _ = temporal_encode(
          x,
          self.temporal_encoding_config,
          self.patches,
          self.hidden_size,
          return_1d=False,
      )
      bs, t, h, w, c = x.shape
      x = x.reshape(bs, t, h * w, c)
  
      def vit_body(
          x, mlp_dim, num_layers, num_heads, encoder_name="Transformer"
      ):
          # If we want to add a class token, add it here.
          if self.classifier in ["token"]:
              n, _, c = x.shape
              cls = self.param(
                  f"cls_{encoder_name}",
                  nn.initializers.zeros,
                  (1, 1, c),
                  x.dtype,
              )
              cls = jnp.tile(cls, [n, 1, 1])
              x = jnp.concatenate([cls, x], axis=1)
  
          x = Encoder(
              temporal_dims=None,  # This is unused for Factorised-Encoder
              mlp_dim=mlp_dim,
              num_layers=num_layers,
              num_heads=num_heads,
              attention_config=self.attention_config,
              dropout_rate=self.dropout_rate,
              attention_dropout_rate=self.attention_dropout_rate,
              stochastic_droplayer_rate=self.stochastic_droplayer_rate,
              dtype=self.dtype,
              name=encoder_name,
          )(x, train=train)
  
          if self.classifier in ["token", "0"]:
              x = x[:, 0]
          elif self.classifier in ("gap", "gmp", "gsp"):
              fn = {"gap": jnp.mean, "gmp": jnp.max, "gsp": jnp.sum}[
                  self.classifier
              ]
              x = fn(x, axis=list(range(1, x.ndim - 1)))
          return x
  
      # run attention across spacec, per frame
      x = jax.vmap(
          functools.partial(
              vit_body,
              mlp_dim=self.spatial_mlp_dim,
              num_layers=self.spatial_num_layers,
              num_heads=self.spatial_num_heads,
              encoder_name="SpatialTransformer",
          ),
          in_axes=1,
          out_axes=1,
          axis_name="time",
      )(x)
      assert x.ndim == 3 and x.shape[:2] == (bs, t)
  
      # run attention across time, over all frames
      if not self.attention_config.get("spatial_only_baseline", False):
          x = vit_body(
              x,
              mlp_dim=self.temporal_mlp_dim,
              num_layers=self.temporal_num_layers,
              num_heads=self.temporal_num_heads,
              encoder_name="TemporalTransformer",
          )
      else:
          # Do global average pooling instead, as method of combining temporal info.
          x = jnp.mean(x, axis=list(range(1, x.ndim - 1)))
  
      if self.representation_size is not None:
          x = nn.Dense(self.representation_size, name="pre_logits")(x)
          x = nn.tanh(x)
      else:
          x = nn_layers.IdentityLayer(name="pre_logits")(x)
  
      if self.return_prelogits:
          return x
      else:
          x = nn.Dense(
              self.num_classes,
              kernel_init=nn.initializers.zeros,
              name="output_projection",
          )(x)
          return x
  
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 SpaceTimeViViT obj's underlying member method code with example usage: obj.__call__(x) and returns: x, x, x