class_docstr: ViViT model with object information.
class_docstr_tok: ["ViViT","model","with","object","information","."]
code: def __call__(
      self,
      inputs: jnp.ndarray,
      boxes: Optional[jnp.ndarray] = None,
      detections: Optional[jnp.ndarray] = None,
      *,
      train: bool = False,
      debug: bool = False,
  ):
      token_score_from_dataloader = self.attach_configs.get(
          "token_score_from_dataloader", False
      )
      data_has_detection = (
          self.detector_configs.get("use_detector", False)
          or (self.attach_configs.get("enabled", False))
          or token_score_from_dataloader
      )
      run_cross_frame_attention = self.attach_configs.get(
          "run_cross_frame_attention", False
      )
      random_object_baseline = self.attach_configs.get(
          "random_object_baseline", False
      )
  
      if data_has_detection:
          assert not random_object_baseline
          if token_score_from_dataloader:
              inputs, token_scores = inputs[..., :3], inputs[..., 3:]
              token_scores = model_utils.resize_token_score(
                  token_scores, self.patches.size
              )  # batch x num_tokens x num_objs
              token_scores = token_scores.transpose(
                  0, 2, 1
              )  #  batch x num_objs x num_tokens
          else:
              assert 0
      elif random_object_baseline:
          n_batch = inputs.shape[0]
          sp = inputs.shape
          sz = self.patches.size
          n_tokens = (sp[1] * sp[2] * sp[3]) // (sz[0] * sz[1] * sz[2])
          token_scores = jax.random.uniform(
              self.make_rng("dropout"), (n_batch, 1, n_tokens)
          )
      else:
          token_scores = None
  
      # Shape is [batch, num_tokens, hidden_size]
      x_tokens, temporal_dims = vivit_model.temporal_encode(
          inputs,
          self.temporal_encoding_config,
          self.patches,
          self.hidden_size,
      )
  
      n_batch, n_tokens, hidden_dim = x_tokens.shape
      height = width = int(np.sqrt(n_tokens // temporal_dims))
      if height * width * temporal_dims != n_tokens:
          raise ValueError("Input is assumed to be square.")
      num_tokens_per_frame = n_tokens // (
          inputs.shape[1] // self.patches.size[2]
      )
  
      #  fg_inds is used to drop background tokens or attach more tokens.
      #  We need to process fg_inds (index of foreground tokens) outside of
      #    CustomEncoder to handle self.classifier == 'token'.
      fg_inds = None
      if self.attach_configs.get(
          "enabled", False
      ) or self.attach_configs.get("drop_pixel_tokens", False):
          fg_inds = model_utils.get_object_inds(
              token_scores, num_tokens_per_frame, self.attach_configs
          )  # batch x num_attach_tokens
  
      # Add positional encodings.
      input_shape = None
      if self.positional_embedding not in ["learned_1d"]:
          x_tokens = model_utils.add_positional_embeddings(
              x_tokens,
              self.positional_embedding,
              input_shape=input_shape,
          )
  
      # If we want to add a class token, add it here.
      if self.classifier == "token":
          cls = self.param(
              "cls",
              nn.initializers.zeros,
              (1, 1, hidden_dim),
              x_tokens.dtype,
          )
          cls = jnp.tile(cls, [n_batch, 1, 1])
          x_tokens = jnp.concatenate([cls, x_tokens], axis=1)
          if fg_inds is not None:
              cls_inds = jnp.zeros((n_batch, 1), dtype=fg_inds.dtype)
              fg_inds = jnp.concatenate(
                  [cls_inds, fg_inds + 1], axis=-1
              )
  
      x_tokens_encoded, _, aux = model_utils.CustomEncoder(
          temporal_dims=temporal_dims,
          hidden_size=self.hidden_size,
          mlp_dim=self.mlp_dim,
          num_layers=self.num_layers,
          num_heads=self.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,
          positional_embedding="none"
          if self.positional_embedding not in ["learned_1d"]
          else self.positional_embedding,
          normalise_output=self.normalise_encoder_output,
          use_approximate_gelu=self.use_approximate_gelu,
          num_tokens_per_frame=num_tokens_per_frame,
          object_config=self.object_config,
          attach_configs=self.attach_configs,
          run_cross_frame_attention=run_cross_frame_attention,
          video_batch=n_batch,
          name="Transformer",
      )(
          x_tokens,
          fg_inds=fg_inds,
          token_scores=token_scores,
          boxes=boxes,
          train=train,
          debug=debug,
      )
  
      if self.classifier in ("token", "0"):
          x_tokens_encoded = x_tokens_encoded[:, 0]
      elif self.classifier in ("gap", "gmp", "gsp"):
          fn = {"gap": jnp.mean, "gmp": jnp.max, "gsp": jnp.sum}[
              self.classifier
          ]
          x_tokens_encoded = fn(x_tokens_encoded, axis=1)
          x_tokens_encoded = nn.LayerNorm(name="encoder_norm")(
              x_tokens_encoded
          )
      else:
          raise ValueError(f"Unknown classifier {self.classifier}")
  
      if self.representation_size is not None:
          x_tokens_encoded = nn.Dense(
              self.representation_size, name="pre_logits"
          )(x_tokens_encoded)
          x_tokens_encoded = nn.tanh(x_tokens_encoded)
      else:
          x_tokens_encoded = nn_layers.IdentityLayer(name="pre_logits")(
              x_tokens_encoded
          )
  
      x_tokens_encoded = nn.Dense(
          self.num_classes,
          kernel_init=nn.initializers.zeros,
          name="output_projection",
      )(x_tokens_encoded)
      if debug:
          return x_tokens_encoded, aux
      return x_tokens_encoded
  
imports: from typing import Any, Optional
  from absl import logging
  import flax
  import flax.linen
  import jax
  import jax.numpy
  import ml_collections
  import numpy
  from scenic.model_lib.base_models import classification_model
  from scenic.model_lib.layers import nn_layers
  from scenic.projects.objectvivit import model_utils
  from scenic.projects.vivit import model
  from scenic.projects.vivit import model_utils
label_desc: Use image node_id 1 for calling the ObjectViViT obj's underlying member method code with example usage: obj.__call__(inputs, boxes, detections) and returns: x_tokens_encoded, x_tokens_encoded,  aux