You can take a look at the HugginFace impl in different TF models.
E.g. TFBert

  

      github.com
  

  
    huggingface/transformers/blob/master/src/transformers/models/bart/modeling_tf_bart.py#L755:L762 1



      
      
          # The tf.debugging asserts are not compliant with XLA then they
          # have to be disabled in other modes than eager.
          if inputs["head_mask"] is not None and tf.executing_eagerly():
              tf.debugging.assert_equal(
                  shape_list(inputs["head_mask"])[0],
                  len(self.layers),
                  message=f"The head_mask should be specified for {len(self.layers)} layers, but it is for {shape_list(inputs['head_mask'])[0]}.",
              )
          
          
# encoder layers
          for idx, encoder_layer in enumerate(self.layers):
          
          
    if inputs["output_hidden_states"]:
                  encoder_states = encoder_states + (hidden_states,)
              # add LayerDrop (see https://arxiv.org/abs/1909.11556 for description)
              dropout_probability = random.uniform(0, 1)
              if inputs["training"] and (dropout_probability < self.layerdrop):  # skip the layer
                  continue
          
          
    hidden_states, attn = encoder_layer(
                  hidden_states,