class_docstr: Pyramid Transformer Encoder.
  
  Attributes:
    mlp_dim: Dimension of the MLP on top of attention block.
    num_layers: Number of layers.
    attention_configs: Configurations passed to the self-attention.
    attention_fn: Self-attention function used in the model.
    transformer_encoder_configs: Transformer encoder configurations.
    dropout_rate:  Dropout rate.
    attention_dropout_rate: Dropout for attention heads.
    classifier: Type of classifier used.
class_docstr_tok: ["Pyramid","Transformer","Encoder",".","Attributes",":","mlp_dim",":","Dimension","of","the","MLP","on","top","of","attention","block",".","num_layers",":","Number","of","layers",".","attention_configs",":","Configurations","passed","to","the","self-attention",".","attention_fn",":","Self-attention","function","used","in","the","model",".","transformer_encoder_configs",":","Transformer","encoder","configurations",".","dropout_rate",":","Dropout","rate",".","attention_dropout_rate",":","Dropout","for","attention","heads",".","classifier",":","Type","of","classifier","used","."]
docstr: Applies Transformer model on the inputs in Pyramid style.
  
  Args:
    inputs: Input data.
    train: If it is training.
  
  Returns:
    Output of a transformer encoder.
docstr_tok: ["Applies","Transformer","model","on","the","inputs","in","Pyramid","style",".","Args",":","inputs",":","Input","data",".","train",":","If","it","is","training",".","Returns",":","Output","of","a","transformer","encoder","."]
code: def __call__(
      self,
      inputs: jnp.ndarray,
      train: bool,
  ) -> jnp.ndarray:
      """Applies Transformer model on the inputs in Pyramid style.
  
      Args:
        inputs: Input data.
        train: If it is training.
  
      Returns:
        Output of a transformer encoder.
      """
      assert inputs.ndim == 3  # Shape is `[batch, len, emb]`.
      x = attention_layers.Add1DPositionEmbedding(
          posemb_init=nn.initializers.normal(stddev=0.02),  # From BERT.
          name="posembed_input",
      )(inputs)
      x = nn.Dropout(rate=self.dropout_rate)(x, deterministic=not train)
  
      pyramid_lens = self.transformer_encoder_configs["pyramid_lens"]
      pyramid_mode = self.transformer_encoder_configs["pyramid_mode"]
      pyramid_attn_func = self.transformer_encoder_configs.get(
          "pyramid_attn_func", []
      )
  
      if len(pyramid_lens) != self.num_layers:
          raise ValueError("Does not support pyramid lens!=num_layers!")
  
      # Input encoder.
      for lyr in range(self.num_layers):
          if not pyramid_attn_func:
              layer_attn_func = self.attention_fn
          else:
              assert len(pyramid_attn_func) == self.num_layers
              layer_attn_func = pyramid_attn_func[lyr]
  
          # TODO(yitay) Support upsampling.
          target_down_scale = abs(pyramid_lens[lyr])
          if (
              target_down_scale != 1
          ):  # Only scale if scaling is more than 1.
              if pyramid_mode == "pool":
                  length = x.shape[1]
                  assert length % target_down_scale == 0, (
                      f"Current length ({length}) must be divisible by scale "
                      f"({target_down_scale})"
                  )
                  x = nn.max_pool(
                      x,
                      (target_down_scale,),
                      strides=(target_down_scale,),
                      padding="VALID",
                  )
  
              elif pyramid_mode == "2d_pool":
                  # Here `-n` means we down-sample by (n, n) on the grid.
                  if self.classifier == "token":
                      x_cls_token, x = x[:, :1, :], x[:, 1:, :]
                  else:
                      x_cls_token, x = x[:, :0, :], x[:, 0:, :]
  
                  bs, l, d = x.shape
                  window_size = int(np.sqrt(l))
                  assert window_size * window_size == l
                  x = x.reshape(bs, window_size, window_size, d)
                  assert window_size % target_down_scale == 0, (
                      f"Current window size ({window_size}) must be divisible by scale "
                      f"({target_down_scale})"
                  )
  
                  x = nn.max_pool(
                      x,
                      (target_down_scale, target_down_scale),
                      strides=(target_down_scale, target_down_scale),
                      padding="VALID",
                  )
                  x = x.reshape(bs, -1, d)
                  x = jnp.concatenate([x_cls_token, x], axis=1)
  
              elif pyramid_mode == "cnn":
                  # Here `-n` means we down-sample by (n, n) on the grid
                  if self.classifier == "token":
                      x_cls_token, x = x[:, :1, :], x[:, 1:, :]
                  else:
                      x_cls_token, x = x[:, :0, :], x[:, 0:, :]
  
                  bs, l, d = x.shape
                  window_size = int(np.sqrt(l))
                  assert window_size * window_size == l
                  x = x.reshape(bs, window_size, window_size, d)
                  assert window_size % target_down_scale == 0, (
                      f"Current window size ({window_size}) must be divisible by scale "
                      f"({target_down_scale})"
                  )
                  x = nn.Conv(
                      x.shape[
                          -1
                      ],  # TODO(dehghani): make this configurable
                      (target_down_scale, target_down_scale),
                      strides=(target_down_scale, target_down_scale),
                      padding="VALID",
                  )(x)
                  x = x.reshape(bs, -1, d)
                  x = jnp.concatenate([x_cls_token, x], axis=1)
  
              elif pyramid_mode == "memory":
                  # TODO(yitay) Inducing set point method in Set Transformer.
                  raise NotImplementedError
  
              elif pyramid_mode == "linformer":
                  # Linformer based downsampling.
                  cur_length = x.shape[-2]
                  assert cur_length % target_down_scale == 0, (
                      f"Current length ({cur_length}) must be divisible by scale "
                      f"({target_down_scale})"
                  )
                  new_length = cur_length // target_down_scale
                  x = nn.LayerNorm(name=f"pool_norm_{lyr}")(x)
                  x = model_utils.LinformerEncoderAttention(
                      num_heads=self.attention_configs["num_heads"],
                      kernel_init=nn.initializers.xavier_uniform(),
                      broadcast_dropout=False,
                      proj_mode="linear",
                      low_rank_features=new_length,
                      downsample=True,
                      dropout_rate=self.attention_dropout_rate,
                  )(x, deterministic=not train)
  
                  # Input encoder.
          x = model_utils.Encoder1DBlock(  # pytype: disable=wrong-arg-types  # jax-ndarray
              mlp_dim=self.mlp_dim,
              attention_fn=layer_attn_func,
              attention_configs=self.attention_configs,
              dropout_rate=self.dropout_rate,
              attention_dropout_rate=self.attention_dropout_rate,
              name=f"encoder_block_{lyr}",
          )(
              x, inputs_kv=None, deterministic=not train
          )
  
      return nn.LayerNorm(name="encoder_norm")(x)
  
imports: from typing import Any, Optional
  import flax.linen
  import jax.numpy
  import ml_collections
  import numpy
  from scenic.model_lib.base_models.multilabel_classification_model import MultiLabelClassificationModel
  from scenic.model_lib.layers import attention_layers
  from scenic.model_lib.layers import nn_layers
  from scenic.model_lib.layers import nn_ops
  from scenic.projects.baselines import vit
  from scenic.projects.fast_vit import model_utils
label_desc: Use image node_id 1 for calling the Encoder1DPyramid obj's underlying member method code with example usage: obj.__call__(inputs, train) and returns: nn