class_docstr: The model can behave as an encoder (with only self-attention) as well as a decoder, in which case a layer of
  cross-attention is added between the self-attention layers, following the architecture described in `Attention is
  all you need <https://arxiv.org/abs/1706.03762>`__ by Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit,
  Llion Jones, Aidan N. Gomez, Lukasz Kaiser and Illia Polosukhin.
class_docstr_tok: ["The","model","can","behave","as","an","encoder","(","with","only","self-attention",")","as","well","as","a","decoder",",","in","which","case","a","layer","of","cross-attention","is","added","between","the","self-attention","layers",",","following","the","architecture","described","in","`","Attention","is","all","you","need","<","https",":","\/\/arxiv.org\/abs\/1706.03762",">","`","__","by","Ashish","Vaswani",",","Noam","Shazeer",",","Niki","Parmar",",","Jakob","Uszkoreit",",","Llion","Jones",",","Aidan","N.","Gomez",",","Lukasz","Kaiser","and","Illia","Polosukhin","."]
code: def __call__(
      self,
      input_ids,
      token_type_ids=None,
      position_ids=None,
      dropout_rng: PRNGKey = None,
      attention_mask=None,
  ):
      (
          input_ids,
          attention_mask,
          token_type_ids,
          position_ids,
      ) = self._check_inputs(
          input_ids, attention_mask, token_type_ids, position_ids
      )
  
      # Handle any PRNG if needed
      rngs = {}
      if dropout_rng is not None:
          rngs["dropout"] = dropout_rng
  
      return self.module.apply(
          {"params": self.params},
          jnp.array(input_ids, dtype="i4"),
          jnp.array(token_type_ids, dtype="i4"),
          jnp.array(position_ids, dtype="i4"),
          jnp.array(attention_mask, dtype="i4"),
          rng=rngs,
      )
  
imports: from typing import Callable, Dict, Tuple
  import numpy
  import flax.linen
  import jax
  import jax.numpy
  from jax.random import PRNGKey
  from modeling_flax_performer_utils import make_fast_softmax_attention
  from transformers.file_utils import add_start_docstrings
  from transformers.modeling_flax_utils import ACT2FN
  from transformers.models.bert.configuration_bert import BertConfig
  from transformers.models.bert.modeling_flax_bert import FlaxBertOnlyMLMHead, FlaxBertPreTrainedModel
  from transformers.utils import logging
label_desc: Use image node_id 4 for calling the FlaxPerformerModel obj's underlying member method code with example usage: obj.__call__(input_ids, token_type_ids, position_ids, dropout_rng, attention_mask) and returns: self