class_docstr: Encoder that supports dropped tokens and object-aware attention.
docstr: Applies Transformer model on the inputs.
  
  Args:
    inputs: array in shape batch x len x emb
    empty_mask: bool array in shape batch x len: whether to keep a token
    fg_inds: int array in shape batch x K, K is smaller then
      len, the index of kept tokens. One of empty_mask and fg_inds
      should be None.
    token_scores: float: batch x num_objects x num_tokens
    boxes: float: batch x T x num_objs x 4 in range 0 - 1
    train: in training or evaluation.
    debug: print debug information
  Returns:
    an array: the updated feature.
docstr_tok: ["Applies","Transformer","model","on","the","inputs",".","Args",":","inputs",":","array","in","shape","batch","x","len","x","emb","empty_mask",":","bool","array","in","shape","batch","x","len",":","whether","to","keep","a","token","fg_inds",":","int","array","in","shape","batch","x","K",",","K","is","smaller","then","len",",","the","index","of","kept","tokens",".","One","of","empty_mask","and","fg_inds","should","be","None",".","token_scores",":","float",":","batch","x","num_objects","x","num_tokens","boxes",":","float",":","batch","x","T","x","num_objs","x","4","in","range","0","-","1","train",":","in","training","or","evaluation",".","debug",":","print","debug","information","Returns",":","an","array",":","the","updated","feature","."]
code: def __call__(
      self,
      inputs: jnp.ndarray,
      empty_mask=None,
      fg_inds=None,
      token_scores=None,
      boxes=None,
      *,
      train: bool = False,
      debug: bool = False,
  ):
      """Applies Transformer model on the inputs.
  
      Args:
        inputs: array in shape batch x len x emb
        empty_mask: bool array in shape batch x len: whether to keep a token
        fg_inds: int array in shape batch x K, K is smaller then
          len, the index of kept tokens. One of empty_mask and fg_inds
          should be None.
        token_scores: float: batch x num_objects x num_tokens
        boxes: float: batch x T x num_objs x 4 in range 0 - 1
        train: in training or evaluation.
        debug: print debug information
      Returns:
        an array: the updated feature.
      """
      assert inputs.ndim == 3  # (batch, len, emb)
      dtype = jax.dtypes.canonicalize_dtype(self.dtype)
      learn_token = self.learn_token_configs.get("enabled", False)
      learn_token_idx = self.learn_token_configs.get("layer_index", 0)
      assert not learn_token
      assert not learn_token_idx
      attach_tokens = self.attach_configs.get("enabled", False)
      drop_pixel_tokens = self.attach_configs.get(
          "drop_pixel_tokens", False
      )
      assert not (attach_tokens and drop_pixel_tokens)
      add_context_tokens = self.attach_configs.get(
          "add_context_tokens", -1
      )
      object_block_idx = self.attach_configs.get("object_block_idx", [])
      drop_block_idx = self.attach_configs.get("drop_block_idx", 0)
  
      n_posembed = self.n_posembed or inputs.shape[1]
      assert (
          n_posembed <= inputs.shape[1]
      ), f"{n_posembed} > {inputs.shape[1]}"
  
      if self.positional_embedding == "sinusoidal_1d":
          x = attention_layers.Add1DPositionEmbedding(posemb_init=None)(
              inputs
          )
      elif self.positional_embedding == "none":
          x = inputs
      else:
          raise ValueError(
              f"Unknown positional embedding {self.positional_embedding}"
          )
      x = nn.Dropout(rate=self.dropout_rate)(x, deterministic=not train)
  
      learned_token_scores = None
      aux = {}
  
      # Input Encoder
      for lyr in range(self.num_layers):
          droplayer_p = (
              lyr / max(self.num_layers - 1, 1)
          ) * self.stochastic_droplayer_rate
  
          if (
              lyr in object_block_idx
              and self.name != "TemporalTransformer"
          ):
              assert lyr > 0, lyr
              x = ObjectBlock(
                  mlp_dim=self.mlp_dim,
                  num_heads=self.num_heads,
                  dropout_rate=self.dropout_rate,
                  attention_dropout_rate=self.attention_dropout_rate,
                  droplayer_p=droplayer_p,
                  name=f"encoderblock_{lyr}",
                  use_approximate_gelu=self.use_approximate_gelu,
                  configs=self.attach_configs,
                  dtype=dtype,
              )(x, token_scores=token_scores, deterministic=not train)
              continue
  
          block = CustomEncoderBlock(
              mlp_dim=self.mlp_dim,
              num_heads=self.num_heads,
              dropout_rate=self.dropout_rate,
              attention_dropout_rate=self.attention_dropout_rate,
              droplayer_p=droplayer_p,
              name=f"encoderblock_{lyr}",
              use_approximate_gelu=self.use_approximate_gelu,
              dtype=dtype,
          )
  
          if (
              drop_pixel_tokens
              and lyr == drop_block_idx
              and fg_inds is not None
          ):
              if add_context_tokens > 0:
                  context_tokens = self._add_context_tokens(
                      x, add_context_tokens, fg_inds
                  )  # batch x num_context_token x hidden_dim
              x, token_scores = self._drop_tokens(
                  x, fg_inds, token_scores, drop_scores=object_block_idx
              )
              # x: batch x num_fg_tokens x hidden_dim
              # token_scores: batch x num_objects x num_fg_tokens
              if add_context_tokens > 0:
                  x = jnp.concatenate([x, context_tokens], axis=1)
                  if object_block_idx and token_scores is not None:
                      num_objs = token_scores.shape[1]
                      num_bg_tokens = context_tokens.shape[1]
                      token_scores = jnp.concatenate(
                          [
                              token_scores,
                              jnp.zeros(
                                  (x.shape[0], num_objs, num_bg_tokens),
                                  jnp.float32,
                              ),
                          ],
                          axis=2,
                      )
  
          x = block(
              x, empty_mask=empty_mask, deterministic=not train
          )  # empty_mask is used here
  
      if self.normalise_output:
          encoded = nn.LayerNorm(name="encoder_norm")(x)
      else:
          encoded = x
  
      return encoded, learned_token_scores, aux
  
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
  from scenic.projects.objectvivit.object_attention import ObjectBlock
label_desc: Use image node_id 1 for calling the CustomEncoder obj's underlying member method code with example usage: obj.__call__(inputs, empty_mask, fg_inds, token_scores, boxes) and returns: encoded,  learned_token_scores,  aux