class_docstr: A simple pyramid transformer encoder.
  
  This structure is inspired by ViTDet (https://arxiv.org/abs/2203.16527). The
  feature pyramid is built using the output from the last layer in the encoder
  and downsampling is performed via strided depthwise convolution. This simple
  design allows us to share the same architecture as the one used in
  classification task.
docstr: Applies Transformer model on the inputs.
code: def __call__(
      self,
      inputs: jnp.ndarray,
      input_mask: Optional[jnp.ndarray] = None,
      train: bool = False,
  ):
      """Applies Transformer model on the inputs."""
  
      num_pyramid_levels = len(
          self.feature_pyramid_config.feature_pyramid_levels
      )
      assert self.num_layers >= num_pyramid_levels
      assert inputs.ndim == 3  # Shape is `[batch, len, emb]`.
      dtype = jax.dtypes.canonicalize_dtype(self.dtype)
      x = self._add_positional_embedding(inputs)
      x = nn.Dropout(rate=self.dropout_rate)(x, deterministic=not train)
  
      num_text_tokens = (
          inputs.shape[1]
          - self.feature_pyramid_config.num_features_level0
      )
      if input_mask is not None:
          input_mask_per_level = model_utils.create_pyramid_input_masks(
              input_mask,
              num_features_level0=self.feature_pyramid_config.num_features_level0,
              num_pyramid_levels=num_pyramid_levels,
              feature_pyramid_downsample_stride=(
                  self.feature_pyramid_config.feature_pyramid_downsample_stride
              ),
              num_text_tokens=num_text_tokens,
          )
          # No downsampling is done in the encoder.
          input_mask = input_mask_per_level[0]
  
      for lyr in range(self.num_layers):
          x = encoders.TransformerEncoder1DBlock(
              mlp_dim=self.mlp_dim,
              num_heads=self.num_heads,
              dropout_rate=self.dropout_rate,
              attention_dropout_rate=self.attention_dropout_rate,
              stochastic_depth=(lyr / max(self.num_layers - 1, 1))
              * self.stochastic_depth,
              name=f"encoderblock_{lyr}",
              dtype=dtype,
          )(
              x,
              input_mask=input_mask,
              deterministic=not train,
          )
      if num_text_tokens:
          text_tokens = x[:, -num_text_tokens:]
      else:
          # No text tokens.
          text_tokens = x[:, :0]
      # Convolution only applies to video tokens.
      x = x[:, : self.feature_pyramid_config.num_features_level0]
      xs = []
      for lyr in range(num_pyramid_levels):
          stride = (
              1
              if lyr == 0
              else self.feature_pyramid_config.feature_pyramid_downsample_stride
          )
          x = nn.Conv(
              features=x.shape[-1],
              kernel_size=(3,),
              strides=(stride,),
              feature_group_count=x.shape[-1],  # Depthwise convolution
              name=f"output_conv_{lyr}",
          )(x)
          x = nn.LayerNorm(name=f"output_ln_{lyr}")(x)
          xs.append(x)
      xs.append(text_tokens)
      return jnp.concatenate(xs, axis=1)
  
imports: import functools
  from typing import Callable, List, Optional, Sequence, Tuple
  from absl import logging
  import flax.linen
  import jax
  import jax.numpy
  import ml_collections
  import numpy
  from scenic.model_lib.layers import attention_layers
  from scenic.projects.baselines import vit
  from scenic.projects.unloc import encoders
  from scenic.projects.unloc import model_utils
label_desc: Use image node_id 1 for calling the SimplePyramidEncoder obj's underlying member method code with example usage: obj.__call__(inputs, input_mask, train) and returns: jnp