class_docstr: Text Transformer.
  
  Attributes:
    vocab_size: Size of the vocabulary.
    features: Number of features.
    num_layers: Number of transformer blocks (self-attn + MLP).
    num_heads: Number of attention heads.
    out_features: Size of the final text embedding.
    use_underscore_module_name: Optionally replace '.' with '_' in parameter
      naming for PAX checkpoint loading.
class_docstr_tok: ["Text","Transformer",".","Attributes",":","vocab_size",":","Size","of","the","vocabulary",".","features",":","Number","of","features",".","num_layers",":","Number","of","transformer","blocks","(","self-attn","+","MLP",")",".","num_heads",":","Number","of","attention","heads",".","out_features",":","Size","of","the","final","text","embedding",".","use_underscore_module_name",":","Optionally","replace","'",".","'","with","'","_","'","in","parameter","naming","for","PAX","checkpoint","loading","."]
code: def __call__(self, text: jnp.ndarray) -> jnp.ndarray:
      positional_embedding = self.param(
          "positional_embedding",
          jax.nn.initializers.zeros,
          (text.shape[1], self.features),
      )
      mask = nn.combine_masks(
          nn.make_attention_mask(text > 0, text > 0),
          nn.make_causal_mask(text),
      )
      x = nn.Embed(
          self.vocab_size, self.features, name="token_embedding"
      )(text)
      x = x + positional_embedding[None]
      x = Transformer(
          self.features,
          self.num_layers,
          self.num_heads,
          use_underscore_module_name=self.use_underscore_module_name,
          name="transformer",
      )(x, attn_mask=mask)
      x = LayerNorm(name="ln_final")(x)
      x = x[jnp.arange(x.shape[0]), text.argmax(-1)]
      x = nn.Dense(
          self.out_features, use_bias=False, name="text_projection"
      )(x)
      return x
  
imports: import functools
  from typing import Sequence, Optional, Union, Tuple
  import flax.linen
  import jax
  import jax.numpy
label_desc: Use image node_id 1 for calling the TextEncoder obj's underlying member method code with example usage: obj.__call__(text) and returns: x