class_docstr: The same as vivit block. Supports object attention and masked tokens.
class_docstr_tok: ["The","same","as","vivit","block",".","Supports","object","attention","and","masked","tokens","."]
docstr: Applies Encoder1DBlock module.
  
  Args:
    inputs: batch x num_fg_tokens x hidden_dim;
    deterministic: bool
    token_scores: optional: batch x num_objs x num_fg_tokens
  Returns:
    x: batch x num_fg_tokens x hidden_dim
docstr_tok: ["Applies","Encoder1DBlock","module",".","Args",":","inputs",":","batch","x","num_fg_tokens","x","hidden_dim",";","deterministic",":","bool","token_scores",":","optional",":","batch","x","num_objs","x","num_fg_tokens","Returns",":","x",":","batch","x","num_fg_tokens","x","hidden_dim"]
code: def __call__(
      self,
      inputs: jnp.ndarray,
      deterministic: bool,
      token_scores: jnp.ndarray,
  ) -> jnp.ndarray:
      """Applies Encoder1DBlock module.
  
      Args:
        inputs: batch x num_fg_tokens x hidden_dim;
        deterministic: bool
        token_scores: optional: batch x num_objs x num_fg_tokens
      Returns:
        x: batch x num_fg_tokens x hidden_dim
      """
      num_tokens_per_object = self.configs.get(
          "num_tokens_per_object", 8
      )
      norm_objects = self.configs.get("norm_objects", False)
      with_objects_linear = self.configs.get(
          "with_objects_linear", True
      )
      with_traj_emb = self.configs.get("with_traj_emb", True)
  
      batch, num_fg_tokens, hidden_dim = inputs.shape
      num_objs = token_scores.shape[1]
  
      patch_tokens = inputs
      object_tokens = patch_tokens[:, None] * token_scores[..., None]
      # batch x num_objs x N x D
  
      if norm_objects:
          object_tokens = object_tokens.sum(axis=2) / (
              token_scores[..., None].sum(axis=2) + 1e-4
          )
          object_tokens = object_tokens[:, :, None, :]
          # batch x num_objs x 1 x D
  
      if with_objects_linear:
          object_tokens = nn.Dense(hidden_dim // 2, use_bias=False)(
              object_tokens
          )
          object_tokens = nn.relu(object_tokens)
          object_tokens = nn.Dense(hidden_dim, use_bias=False)(
              object_tokens
          )
          object_tokens = nn.relu(
              object_tokens
          )  # batch x num_objs x N x D
  
      if not norm_objects:
          object_tokens = object_tokens.transpose(
              0, 1, 3, 2
          )  # batch x num_objs x D x N
          object_tokens, _ = jax.lax.top_k(
              object_tokens, k=num_tokens_per_object
          )  # batch x num_objs x D x k
          object_tokens = object_tokens.transpose(
              0, 1, 3, 2
          )  # batch x num_objs x k x D
  
      if with_traj_emb:
          box_categories = self.param(
              "box_categories",
              nn.initializers.zeros,
              (1, num_objs, 1, hidden_dim),
              jnp.float32,
          )
          object_tokens = object_tokens + box_categories
  
      object_tokens = object_tokens.reshape(batch, -1, hidden_dim)
      all_tokens = jnp.concatenate(
          [patch_tokens, object_tokens],
          axis=1,
      )  # batch x (num_fg_tokens + num_objs * k) x D
  
      # Attention block.
      x = nn.LayerNorm(dtype=self.dtype)(all_tokens)
      x = nn.MultiHeadDotProductAttention(
          num_heads=self.num_heads,
          kernel_init=self.attention_kernel_initializer,
          broadcast_dropout=False,
          dropout_rate=self.attention_dropout_rate,
          attention_fn=self.attention_fn,
          dtype=self.dtype,
      )(x, x, deterministic=deterministic)
      x = nn.Dropout(rate=self.dropout_rate)(x, deterministic)
  
      x = x[:, :num_fg_tokens, :]
  
      drop_pattern = self.get_drop_pattern(x, deterministic)
      x = x * (1.0 - drop_pattern) + inputs
  
      # MLP block.
      y = nn.LayerNorm(dtype=self.dtype)(x)
      y = attention_layers.MlpBlock(  # pytype: disable=wrong-arg-types  # jnp-type
          mlp_dim=self.mlp_dim,
          dtype=self.dtype,
          dropout_rate=self.dropout_rate,
          activation_fn=functools.partial(
              nn.gelu, approximate=self.use_approximate_gelu
          ),
          kernel_init=self.mlp_kernel_initializer,
          bias_init=self.mlp_bias_initializer,
      )(
          y, deterministic=deterministic
      )
  
      drop_pattern = self.get_drop_pattern(x, deterministic)
      ret = y * (1.0 - drop_pattern) + x
      return ret
  
file_path: scenic/scenic/projects/objectvivit/object_attention.py
imports: import functools
  from typing import Any, Callable, Iterable, Optional
  import flax.linen
  import jax
  import jax.numpy
  import ml_collections
  from scenic.model_lib.layers import attention_layers
label_desc: Use image node_id 2 for calling the ObjectBlock obj's underlying member method code with example usage: obj.__call__(inputs, deterministic, token_scores) and returns: ret