class_docstr: Stores N BertLayer(s)
class_docstr_tok: ["Stores","N","BertLayer","(","s",")"]
code: def __call__(self, inputs, attention_mask):
      assert (
          self.num_layers > 0
      ), f"num_layers should be >= 1, got ({self.num_layers})"
  
      # Initialize input / output
      input_i = inputs
  
      # Forward over all encoders
      for i in range(self.num_layers):
          layer = FlaxPerformerLayer(
              self.num_heads,
              self.head_size,
              self.intermediate_size,
              hidden_act=self.hidden_act,
              name=f"{i}",
          )
          input_i = layer(input_i, attention_mask)
      return input_i
  
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 1 for calling the FlaxPerformerLayerCollection obj's underlying member method code with example usage: obj.__call__(inputs, attention_mask) and returns: input_i