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 convert_from_pytorch(pt_state: Dict, config: BertConfig) -> Dict:
      jax_state = dict(pt_state)
  
      # Need to change some parameters name to match Flax names so that we don't have to fork any layer
      for key, tensor in pt_state.items():
          # Key parts
          key_parts = set(key.split("."))
  
          # Every dense layer has "kernel" parameters instead of "weight"
          if "dense.weight" in key:
              del jax_state[key]
              key = key.replace("weight", "kernel")
              jax_state[key] = tensor
  
          # SelfAttention needs also to replace "weight" by "kernel"
          if {"query", "key", "value"} & key_parts:
              # Flax SelfAttention decomposes the heads (num_head, size // num_heads)
              if "bias" in key:
                  jax_state[key] = tensor.reshape(
                      (config.num_attention_heads, -1)
                  )
              elif "weight":
                  del jax_state[key]
                  key = key.replace("weight", "kernel")
                  tensor = tensor.reshape(
                      (
                          config.num_attention_heads,
                          -1,
                          config.hidden_size,
                      )
                  ).transpose((2, 0, 1))
                  jax_state[key] = tensor
  
          # SelfAttention output is not a separate layer, remove one nesting
          if "attention.output.dense" in key:
              del jax_state[key]
              key = key.replace(
                  "attention.output.dense", "attention.self.out"
              )
              jax_state[key] = tensor
  
          # SelfAttention output is not a separate layer, remove nesting on layer norm
          if "attention.output.LayerNorm" in key:
              del jax_state[key]
              key = key.replace(
                  "attention.output.LayerNorm", "attention.LayerNorm"
              )
              jax_state[key] = tensor
  
          # There are some transposed parameters w.r.t their PyTorch counterpart
          if (
              "intermediate.dense.kernel" in key
              or "output.dense.kernel" in key
          ):
              jax_state[key] = tensor.T
  
          # Self Attention output projection needs to be transposed
          if "out.kernel" in key:
              jax_state[key] = tensor.reshape(
                  (config.hidden_size, config.num_attention_heads, -1)
              ).transpose(1, 2, 0)
  
          # Pooler needs to transpose its kernel
          if "pooler.dense.kernel" in key:
              jax_state[key] = tensor.T
  
          # Handle LayerNorm conversion
          if "LayerNorm" in key:
              del jax_state[key]
  
              # Replace LayerNorm by layer_norm
              new_key = key.replace("LayerNorm", "layer_norm")
  
              if "weight" in key:
                  new_key = new_key.replace("weight", "gamma")
              elif "bias" in key:
                  new_key = new_key.replace("bias", "beta")
  
              jax_state[new_key] = tensor
  
      return jax_state
  
imports: from typing import Callable, Dict, Tuple
  import flax.linen
  import jax
  import jax.numpy
  import 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 1 for calling the FlaxPerformerModel obj's underlying member method code with example usage: obj.convert_from_pytorch(pt_state, config) and returns: jax_state