code: def test_layer_stack_call_num_unique_layers(self):
      x = np.random.randn(2, 3, 4, 5)
      layer_stack, context = create_dummy_model(
          self.converter.mesh, shapes=x.shape
      )
      x_mtf = self.converter.convert_np_array_to_mtf_tensor(
          x, dim_names=self.default_dim_names, dtype=tf.float32
      )
      output_mtf = layer_stack.call(context, x_mtf)
      lowering, _ = self.converter.convert_mtf_tensor_to_tf_tensor(
          output_mtf
      )
  
      # Test the number of unique layers.
      all_vars = lowering.graph.all_variables
      self_attn_vars = [
          var.name for var in all_vars if "SelfAttention" in var.name
      ]
  
      # We expect total of `n_layers` of SelfAttention and DenseReluDense layers.
      n_layers = len(layer_stack.layers)
  
      # We expect n_sublayers` SelfAttention.
      n_sublayers = n_layers // 2
  
      # Each self attn layer has 4 variables: wq, wk, wv, wo.
      self.assertEqual(len(self_attn_vars) // 4, n_sublayers)
  
file_path: mesh/mesh_tensorflow/transformer/funnel_transformer_test.py
imports: from absl.testing import parameterized
  import mesh_tensorflow
  from mesh_tensorflow import test_utils
  from mesh_tensorflow.transformer import funnel_transformer
  from mesh_tensorflow.transformer import transformer
  from mesh_tensorflow.transformer import transformer_layers
  import numpy
  import tensorflow.compat.v1
label_desc: Use image node_id 5 for calling the FunnelTransformerTest obj's underlying member method code with example usage: obj.test_layer_stack_call_num_unique_layers() without return types