code: def balance_layers_simple(n_layers: int, n_chunks: int):
      """
      ### Balance layers
  
      Split the `n_layers` into `n_chunks`. This is used for pipeline parallel training.
  
      :param n_layers: is the number of layers
      :param n_chunks: is the number of chunks
      :return: returns a list with the number of layers for each chunk
      """
      balance = []
      for i in range(n_chunks):
          balance.append((n_layers - sum(balance)) // (n_chunks - i))
  
      return list(reversed(balance))
  
file_path: annotated_deep_learning_paper_implementations/labml_nn/neox/utils/__init__.py
imports: import typing
  from typing import List, Optional
  import torch
  from labml import logger
  from labml.logger import Text
  from labml_nn.neox.tokenizer import get_tokenizer
label_desc: Use image node_id 4 for calling a global function with example usage: balance_layers_simple(n_layers, n_chunks) and returns: list