class_docstr: BERT text encoder.
class_docstr_tok: ["BERT","text","encoder","."]
code: def __call__(  # pytype: disable=signature-mismatch  # jax-ndarray
      self,
      text: jnp.ndarray,
      mask: Optional[jnp.ndarray] = None,
      *,
      text_for_mlm: Optional[jnp.ndarray] = None,
      segment_ids_for_mlm: Optional[jnp.ndarray] = None,
      mask_for_mlm: Optional[jnp.ndarray] = None,
      masked_lm_positions: Optional[jnp.ndarray] = None,
      train: bool = False,
      debug: bool = False,
  ) -> Union[jnp.ndarray, tuple[jnp.ndarray, jnp.ndarray]]:
      config = CONFIGS[self.config_name]
      config.dtype = self.dtype
      config.dropout_rate = 0.0
  
      # We leave the default eps because the LayerNorm is always computed in FP32.
      encoder = bert.BertEncoder(**dataclasses.asdict(config))
      mlm_head = bert_heads.MLMHead(
          encoder=encoder,
          hidden_size=config.hidden_size,
          vocab_size=config.vocab_size,
          kernel_init=config.kernel_init,
          dropout_rate=config.dropout_rate,
          dtype=config.dtype,
      )
      pooler = bert_heads.BertPooler(
          kernel_init=config.kernel_init, dtype=config.dtype
      )
      cls_head = bert_heads.ClassifierHead(
          pooler=pooler,
          num_classes=self.embedding_size,
          kernel_init=config.kernel_init,
          dropout_rate=config.dropout_rate,
          enable_dropout=not train,
          dtype=config.dtype,
      )
  
      encoded_tokens = encoder(
          token_ids=text,
          segment_ids=jnp.zeros_like(text),
          position_ids=jnp.arange(text.shape[-1]),
          input_mask=mask,
          enable_dropout=train,
      )
  
      if self.pretraining_mode:
          encoded_sequence = cls_head(encoded_tokens)
          if self.compute_mlm:
              encoded_tokens_for_mlm = encoder(
                  token_ids=text_for_mlm,
                  segment_ids=segment_ids_for_mlm,
                  position_ids=jnp.arange(
                      text_for_mlm.shape[-1]
                  ),  # pytype: disable=attribute-error
                  input_mask=mask_for_mlm,
                  enable_dropout=train,
              )
              mlm_logits = mlm_head(
                  encoded_tokens_for_mlm,
                  masked_positions=masked_lm_positions,
              )
              return encoded_sequence, mlm_logits
          else:
              return encoded_sequence
      else:
          return (
              encoded_tokens
              if self.return_all_tokens
              else encoded_tokens[:, 0]
          )
  
imports: import dataclasses
  from typing import Any
  from typing import Optional
  from typing import Union
  import flax
  from flax.core import frozen_dict
  import flax.linen
  import jax
  import jax.numpy
  from scenic.projects.lang4video.model.base_encoders import TextEncoder
  import tensorflow
  from flaxformer.architectures.bert import bert
  from flaxformer.architectures.bert import bert_checkpoint_converter
  import flaxformer.architectures.bert.configs
  import flaxformer.architectures.bert.heads
label_desc: Use image node_id 1 for calling the BertTextEncoder obj's underlying member method code with example usage: obj.__call__(text, mask) and returns: unknown, encoded_sequence,  mlm_logits, encoded_sequence