class_docstr: PonderNet Transformer Encoder.
  
  Attributes:
    num_layers: Number of layers.
    mlp_dim: Dimension of the mlp on top of attention block.
    inputs_positions: Input subsequence positions for packed examples.
    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 timm
      library, which does per-example layer dropping and uses independent
      dropping patterns for each skip-connection.
    dtype: Dtype of activations.
docstr: Applies Transformer model on the inputs.
docstr_tok: ["Applies","Transformer","model","on","the","inputs","."]
code: def __call__(self, inputs: jnp.ndarray, *, train: bool = False):
      """Applies Transformer model on the inputs."""
      assert inputs.ndim == 3  # Shape is `[batch, len, emb]`.
      x = vit.AddPositionEmbs(
          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)
  
      dtype = jax.dtypes.canonicalize_dtype(self.dtype)
  
      # We use layers.AdaptiveComputationTime only when we are doing ACT.
      if self.ac_config is None:
          # We make the layer first if we are using parameter sharing.
          if not self.parameter_sharing:
              for i in range(self.num_layers):
                  x = Encoder1DBlock(
                      mlp_dim=self.mlp_dim,
                      num_heads=self.num_heads,
                      dropout_rate=self.dropout_rate,
                      attention_dropout_rate=self.attention_dropout_rate,
                      name="encoderblock_" + str(i),
                      dtype=dtype,
                  )(x, deterministic=not train)
              else:
                  encoder_block = Encoder1DBlock(
                      mlp_dim=self.mlp_dim,
                      num_heads=self.num_heads,
                      dropout_rate=self.dropout_rate,
                      attention_dropout_rate=self.attention_dropout_rate,
                      name="encoderblock",
                      dtype=dtype,
                  )
                  for i in range(self.num_layers):
                      x = encoder_block(x, deterministic=not train)
          auxiliary_outputs = None
      else:
          encoder_block = Encoder1DBlock(
              mlp_dim=self.mlp_dim,
              num_heads=self.num_heads,
              dropout_rate=self.dropout_rate,
              attention_dropout_rate=self.attention_dropout_rate,
              name="encoderblock",
              dtype=dtype,
          )
          x, auxiliary_outputs = layers.AdaptiveComputationTime(
              self.ac_config,
              encoder_block,
              self.parameter_sharing,
              name="act",
          )(x, not train)
          (all_states, all_p, n_updates) = auxiliary_outputs
      encoded_norm_layer = nn.LayerNorm(name="encoder_norm")
      encoded = encoded_norm_layer(x)
      all_states = encoded_norm_layer(all_states)
      return encoded, (all_states, all_p, n_updates)
  
imports: from typing import Any, Optional
  import flax.linen
  from flax.training import common_utils
  import jax
  import jax.numpy
  import ml_collections
  import numpy
  from scenic.model_lib.base_models import base_model
  from scenic.model_lib.base_models import model_utils
  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.projects.baselines import vit
  from scenic.projects.baselines.pondernet import layers
label_desc: Use image node_id 1 for calling the PonderNetEncoder obj's underlying member method code with example usage: obj.__call__(inputs) and returns: encoded,  all_states,  all_p,  n_updates