class_docstr: Transformer.
class_docstr_tok: ["Transformer","."]
code: def __call__(
      self,
      x: jnp.ndarray,  # Shape: (N, ..., X)
      mask: Optional[jnp.ndarray] = None,
      *,
      train: bool = False,
      debug: bool = False,
  ) -> jnp.ndarray:  # Shape: (N, ..., E)
      if mask is None:
          encoder_mask = logit_mask = None
      else:
          encoder_mask = dense_attention.make_attention_mask(
              mask, mask, dtype=self.dtype
          )
          logit_mask = mask[..., jnp.newaxis]
  
      for _ in range(self.num_layers):
          x = t5_1_1.t5_common_layers.encoder_layer(
              num_heads=self.num_heads,
              head_dim=self.head_dim,
              mlp_dim=self.mlp_dim,
              dropout_rate=self.dropout_rate,
              activations=self.activations,
              dtype=self.dtype,
          )(
              x,
              encoder_mask=encoder_mask,
              logit_mask=logit_mask,
              enable_dropout=train,
          )
  
      x = layer_norm.T5LayerNorm(dtype=self.dtype)(x)
      x = nn.Dropout(rate=self.dropout_rate, broadcast_dims=(-2,))(
          x, deterministic=not train
      )
  
      if logit_mask is not None:
          x = logit_mask * x
  
      return x if self.return_all_tokens else x[:, 0]
  
imports: from collections.abc import Callable
  from collections.abc import Sequence
  from typing import Optional
  from typing import Union
  import flax.linen
  import jax.numpy
  from scenic.projects.lang4video.model.base_encoders import ImageEncoder
  from scenic.projects.lang4video.model.base_encoders import TextEncoder
  from flaxformer.architectures.t5 import t5_1_1
  from flaxformer.components import layer_norm
  from flaxformer.components.attention import dense_attention
label_desc: Use image node_id 1 for calling the Transformer obj's underlying member method code with example usage: obj.__call__(x, mask) and returns: unknown