code: def create_dummy_model(
      mesh,
      shapes,
      n_blocks=2,
      block_param_size_str="2_2",
      block_repeat_size_str="1_1",
  ):
      """Creates a dummy model and layer stack with 4-dimensional input."""
  
      assert len(shapes) == 4
      outer_batch_size, batch_size, length, d_model = shapes
      batch_dim = mtf.Dimension("batch", batch_size)
      outer_batch_dim = mtf.Dimension("outer_batch", outer_batch_size)
      length_dim = mtf.Dimension("length", length)
      block_param_size = list(map(int, block_param_size_str.split("_")))
      block_repeat_size = list(
          map(int, block_repeat_size_str.split("_"))
      )
  
      sublayers_initial = [
          transformer.sublayer_dropout,
      ]
      sublayers_per_layer = [
          transformer.sublayer_rms_norm,
          transformer.sublayer_call_layer,
          transformer.sublayer_dropout,
          transformer.sublayer_residual,
      ]
      sublayers_final = [
          transformer.sublayer_rms_norm,
          transformer.sublayer_dropout,
      ]
      submodules = [
          transformer_layers.SelfAttention(),
          transformer_layers.DenseReluDense(),
      ]
  
      n_sublayers = np.array(block_param_size).prod()
      layers = submodules * n_sublayers
      layer_stack = funnel_transformer.FunnelTransformerLayerStack(
          layers=layers,
          n_blocks=n_blocks,
          block_param_size=block_param_size,
          block_repeat_size=block_repeat_size,
          sublayers_initial=sublayers_initial,
          sublayers_per_layer=sublayers_per_layer,
          sublayers_final=sublayers_final,
      )
  
      model = transformer.Unitransformer(
          input_vocab_size=10,
          output_vocab_size=10,
          autoregressive=False,
          max_length=8,
          d_model=d_model,
          layer_stack=layer_stack,
      )
  
      context = transformer.Context(
          model=model,
          mesh=mesh,
          batch_dims=[batch_dim, outer_batch_dim],
          length_dim=length_dim,
          variable_dtype=mtf.VariableDType(tf.float32),
          sequence_id=mtf.ones(mesh, mtf.Shape([length_dim])),
          position=mtf.range(mesh, length_dim, dtype=tf.int32),
      )
      return layer_stack, context
  
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 1 for calling a global function with example usage: create_dummy_model(mesh, shapes, n_blocks, block_param_size_str, block_repeat_size_str) and returns: layer_stack,  context