code: def _split_into_groups(n, max_group_size, mesh_dim_size):
      """Helper function for figuring out how to split a dimension into groups.
  
      We have a dimension with size n and we want to split it into
      two dimensions: n = num_groups * group_size
  
      group_size should be the largest possible value meeting the constraints:
        group_size <= max_group_size
        (num_groups = n/group_size) is a multiple of mesh_dim_size
  
      Args:
        n: an integer
        max_group_size: an integer
        mesh_dim_size: an integer
  
      Returns:
        num_groups: an integer
        group_size: an integer
  
      Raises:
        ValueError: if n is not a multiple of mesh_dim_size
      """
      if n % mesh_dim_size != 0:
          raise ValueError(
              "n=%d is not a multiple of mesh_dim_size=%d"
              % (n, mesh_dim_size)
          )
      num_groups = max(1, n // max_group_size)
      while num_groups % mesh_dim_size != 0 or n % num_groups != 0:
          num_groups += 1
      group_size = n // num_groups
      tf.logging.info(
          "_split_into_groups(n=%d, max_group_size=%d, mesh_dim_size=%d)"
          " = (num_groups=%d group_size=%d)"
          % (n, max_group_size, mesh_dim_size, num_groups, group_size)
      )
      return num_groups, group_size
  
imports: from __future__ import absolute_import
  from __future__ import division
  from __future__ import print_function
  import gin
  import mesh_tensorflow
  from mesh_tensorflow.transformer import transformer
  import tensorflow.compat.v1
label_desc: Use image node_id 13 for calling a global function with example usage: _split_into_groups(n, max_group_size, mesh_dim_size) and returns: num_groups,  group_size