class_docstr: Transformer feature pyramid encoder.
  
  Attributes:
    num_layers: Number of layers.
    mlp_dim: Dimension of the mlp on top of attention block.
    num_heads: Number of attention heads.
    feature_pyramid_config: Feature pyramid config.
    downsample_strategy: Strategy to downsample video tokens. Options are:
      'subsample', 'avg_pool', or 'max_pool'.
    positional_embedding: 'learned', 'sinusoid', or 'none'.
    positional_embedding_max_length: If set, the positional embeddings/encodings
      are applied to the first N elements in the input sequence. If not set, the
      positional embeddings/encodings are added to the entire input sequence.
    dropout_rate: Dropout rate.
    stochastic_depth: probability of dropping a layer linearly grows from 0 to
      the provided value. Our implementation of stochastic depth follows the
      timm library, which does per-example layer dropping and uses independent
      dropping patterns for each skip-connection.
    window_size: Window size for window attention blocks.
    window_block_indexes: Tuple. Indexes for blocks using window attention.
    dtype: Dtype of activations.
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,
          )
      else:
          input_mask_per_level = [None] * num_pyramid_levels
  
      cur_pyramid_level = 0
      xs = []
      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_per_level[cur_pyramid_level],
              deterministic=not train,
          )
          if lyr in self.feature_pyramid_config.feature_pyramid_levels:
              # Convolution only applies to video tokens.
              if num_text_tokens:
                  xs.append(x[:, :-num_text_tokens])
              else:
                  xs.append(x)
              cur_pyramid_level += 1
              if (
                  self.feature_pyramid_config.feature_pyramid_downsample_stride
                  > 1
                  and lyr < self.num_layers - 1
              ):
                  sampled_indices = np.arange(
                      0,
                      self.feature_pyramid_config.num_features_level0,
                      self.feature_pyramid_config.feature_pyramid_downsample_stride
                      ** cur_pyramid_level,
                      dtype=np.int32,
                  )
                  x = self._downsample_video_tokens(
                      x,
                      num_text_tokens,
                      sampled_indices=sampled_indices,
                  )
      if num_text_tokens:
          text_tokens = x[:, -num_text_tokens:]
      else:
          # No text tokens.
          text_tokens = x[:, :0]
      xs = self._build_top_down_path(xs)
      xs = self._apply_output_conv(xs)
      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 10 for calling the FeaturePyramidEncoder obj's underlying member method code with example usage: obj.__call__(inputs, input_mask, train) and returns: jnp