class_docstr: The base class for all the policies. For each different model, it should have a different policy class,
  like BertPolicy for Bert Model or OPTPolicy for OPT model.
  
  Shardformer has provided many built-in sharding policies for the mainstream models. You can use the
  built-in policies by setting `policy = None`, which is already the default argument for `Shardformer.optimize`.
  If you want to define your own policy, you can inherit from this class and overwrite the methods you want to modify.
class_docstr_tok: ["The","base","class","for","all","the","policies",".","For","each","different","model",",","it","should","have","a","different","policy","class",",","like","BertPolicy","for","Bert","Model","or","OPTPolicy","for","OPT","model",".","Shardformer","has","provided","many","built-in","sharding","policies","for","the","mainstream","models",".","You","can","use","the","built-in","policies","by","setting","`","policy","=","None","`",",","which","is","already","the","default","argument","for","`","Shardformer.optimize","`",".","If","you","want","to","define","your","own","policy",",","you","can","inherit","from","this","class","and","overwrite","the","methods","you","want","to","modify","."]
docstr: Get the start index and end index of layers for each stage.
  
  Args:
      layers_per_stage (List[int]): number of layers for each stage
      stage (int): the stage index
      num_stages (int): number of stages
      num_model_chunks (int): number of model chunks
  
  Returns:
      - Tuple[int, int]: the start index and end index of this stage
      - List[Tuple[int, int]]: the start index and end index of this stage for each model chunk
docstr_tok: ["Get","the","start","index","and","end","index","of","layers","for","each","stage",".","Args",":","layers_per_stage","(","List","[","int","]",")",":","number","of","layers","for","each","stage","stage","(","int",")",":","the","stage","index","num_stages","(","int",")",":","number","of","stages","num_model_chunks","(","int",")",":","number","of","model","chunks","Returns",":","-","Tuple","[","int",",","int","]",":","the","start","index","and","end","index","of","this","stage","-","List","[","Tuple","[","int",",","int","]","]",":","the","start","index","and","end","index","of","this","stage","for","each","model","chunk"]
code: def get_stage_index(
      layers_per_stage: List[int],
      stage: int,
      num_model_chunks: int = 1,
      num_stages: int = 0,
  ) -> Union[Tuple[int, int], List[Tuple[int, int]]]:
      """
      Get the start index and end index of layers for each stage.
  
      Args:
          layers_per_stage (List[int]): number of layers for each stage
          stage (int): the stage index
          num_stages (int): number of stages
          num_model_chunks (int): number of model chunks
  
      Returns:
          - Tuple[int, int]: the start index and end index of this stage
          - List[Tuple[int, int]]: the start index and end index of this stage for each model chunk
  
      """
      num_layers_per_stage_accumulated = np.insert(
          np.cumsum(layers_per_stage), 0, 0
      )
  
      stage_indices = []
      for model_chunk in range(num_model_chunks):
          start_idx = num_layers_per_stage_accumulated[
              stage + model_chunk * num_stages
          ]
          end_idx = num_layers_per_stage_accumulated[
              stage + model_chunk * num_stages + 1
          ]
          stage_indices.append([start_idx, end_idx])
  
      return (
          stage_indices[0] if num_model_chunks == 1 else stage_indices
      )
  
imports: from abc import ABC, abstractmethod
  from dataclasses import dataclass
  from typing import Any, Callable, Dict, List, Optional, Tuple, Union
  import numpy
  import torch.nn
  from torch import Tensor
  from torch.nn import Module
  from colossalai.pipeline.stage_manager import PipelineStageManager
  from ..layer.normalization import BaseLayerNorm
  from ..layer.parallel_module import ParallelModule
  from ..shard.shard_config import ShardConfig
label_desc: Use image node_id 14 for calling the Policy obj's underlying member method code with example usage: obj.get_stage_index(layers_per_stage, stage, num_model_chunks, num_stages) and returns: unknown