docstr: Input the distribution of layers among stages, the current stage and the first stage of decoder.
  Return the starting/ending idx of layers in encoder/decoder
docstr_tok: ["Input","the","distribution","of","layers","among","stages",",","the","current","stage","and","the","first","stage","of","decoder",".","Return","the","starting\/ending","idx","of","layers","in","encoder\/decoder"]
code: def get_t5_stage_index(
      layers_per_stage: List[int],
      stage: int,
      decoder_starting_stage: int,
  ) -> Tuple[bool, int, int]:
      """
      Input the distribution of layers among stages, the current stage and the first stage of decoder.
      Return the starting/ending idx of layers in encoder/decoder
      """
      if stage < decoder_starting_stage:
          return Policy.get_stage_index(
              layers_per_stage[:decoder_starting_stage], stage
          )
      else:
          return Policy.get_stage_index(
              layers_per_stage[decoder_starting_stage:],
              stage - decoder_starting_stage,
          )
  
imports: import warnings
  from functools import partial
  from typing import Callable, Dict, List, Tuple
  import numpy
  from torch import Tensor, nn
  from colossalai.shardformer.layer import DropoutForParallelInput, Embedding1D, FusedRMSNorm, Linear1D_Col, Linear1D_Row, RMSNorm, VocabParallelEmbedding1D
  from colossalai.shardformer.policies.base_policy import ModulePolicyDescription
  from ..modeling.jit import get_jit_fused_dropout_add_func
  from ..modeling.t5 import T5PipelineForwards, get_jit_fused_T5_layer_ff_forward, get_t5_flash_attention_forward, get_T5_layer_cross_attention_forward, get_T5_layer_self_attention_forward
  from .base_policy import ModulePolicyDescription, Policy, SubModuleReplacementDescription
label_desc: Use image node_id 6 for calling the T5BasePolicy obj's underlying member method code with example usage: obj.get_t5_stage_index(layers_per_stage, stage, decoder_starting_stage) and returns: Policy, Policy