class_docstr: Information to be passed and used for a batch of inputs during
  a single model forward
code: def init_from_batch(
      cls,
      batch: torch.Tensor,
      max_input_len: int,
      max_output_len: int,
      cache_manager: MemoryManager,
  ):
      if not isinstance(
          batch, (BatchEncoding, dict, list, torch.Tensor)
      ):
          raise TypeError(
              f"batch type {type(batch)} is not supported in prepare_batch_state"
          )
  
      input_ids_list = None
      attention_mask = None
  
      if isinstance(batch, (BatchEncoding, dict)):
          input_ids_list = batch["input_ids"]
          attention_mask = batch["attention_mask"]
      else:
          input_ids_list = batch
      if isinstance(input_ids_list[0], int):  # for a single input
          input_ids_list = [input_ids_list]
          attention_mask = (
              [attention_mask]
              if attention_mask is not None
              else attention_mask
          )
  
      batch_size = len(input_ids_list)
  
      seq_start_indexes = torch.zeros(
          batch_size, dtype=torch.int32, device="cuda"
      )
      seq_lengths = torch.zeros(
          batch_size, dtype=torch.int32, device="cuda"
      )
      start_index = 0
  
      max_len_in_batch = -1
      if isinstance(batch, (BatchEncoding, dict)):
          for i, attn_mask in enumerate(attention_mask):
              curr_seq_len = len(attn_mask)
              seq_lengths[i] = curr_seq_len
              seq_start_indexes[i] = start_index
              start_index += curr_seq_len
              max_len_in_batch = (
                  curr_seq_len
                  if curr_seq_len > max_len_in_batch
                  else max_len_in_batch
              )
      else:
          length = max(len(input_id) for input_id in input_ids_list)
          for i, input_ids in enumerate(input_ids_list):
              curr_seq_len = length
              seq_lengths[i] = curr_seq_len
              seq_start_indexes[i] = start_index
              start_index += curr_seq_len
              max_len_in_batch = (
                  curr_seq_len
                  if curr_seq_len > max_len_in_batch
                  else max_len_in_batch
              )
      block_loc = torch.zeros(
          (batch_size, max_input_len + max_output_len),
          dtype=torch.long,
          device="cuda",
      )
  
      return cls(
          batch_size=batch_size,
          max_len_in_batch=max_len_in_batch,
          seq_len=seq_lengths.to("cuda"),
          start_loc=seq_start_indexes.to("cuda"),
          block_loc=block_loc,
          decode_layer_id=0,
          past_key_values_len=0,
          is_context_stage=True,
          cache_manager=cache_manager,
      )
  
file_path: colossalai/colossalai/legacy/inference/tensor_parallel/batch_infer_state.py
imports: from dataclasses import dataclass
  import torch
  from transformers.tokenization_utils_base import BatchEncoding
  from .kvcache_manager import MemoryManager
label_desc: Use image node_id 4 for calling the BatchInferState obj's underlying member method code with example usage: obj.init_from_batch(cls, batch, max_input_len, max_output_len, cache_manager) and returns: cls