class_docstr: Engine class for tensor parallel inference.
  
  Args:
      model (Module): original model, e.g. huggingface CausalLM
      shard_config (ShardConfig): The config for sharding original model
      max_batch_size (int): maximum batch size
      max_input_len (int): maximum input length of sequence
      max_output_len (int): maximum output length of output tokens
      dtype (torch.dtype): datatype used to init KV cache space
      device (str): device the KV cache of engine to be initialized on
  
  Examples:
      >>> # define model and shard config for your inference
      >>> model = ...
      >>> generate_kwargs = ...
      >>> shard_config = ShardConfig(enable_tensor_parallelism=True, extra_kwargs={"inference_only": True})
      >>> infer_engine = TPInferEngine(model, shard_config, MAX_BATCH_SIZE, MAX_INPUT_LEN, MAX_OUTPUT_LEN)
      >>> outputs = infer_engine.generate(input_ids, **generate_kwargs)
class_docstr_tok: ["Engine","class","for","tensor","parallel","inference",".","Args",":","model","(","Module",")",":","original","model",",","e.g",".","huggingface","CausalLM","shard_config","(","ShardConfig",")",":","The","config","for","sharding","original","model","max_batch_size","(","int",")",":","maximum","batch","size","max_input_len","(","int",")",":","maximum","input","length","of","sequence","max_output_len","(","int",")",":","maximum","output","length","of","output","tokens","dtype","(","torch.dtype",")",":","datatype","used","to","init","KV","cache","space","device","(","str",")",":","device","the","KV","cache","of","engine","to","be","initialized","on","Examples",":",">",">",">","#","define","model","and","shard","config","for","your","inference",">",">",">","model","=","...",">",">",">","generate_kwargs","=","...",">",">",">","shard_config","=","ShardConfig","(","enable_tensor_parallelism=True",",","extra_kwargs=","{","``","inference_only","''",":","True","}",")",">",">",">","infer_engine","=","TPInferEngine","(","model",",","shard_config",",","MAX_BATCH_SIZE",",","MAX_INPUT_LEN",",","MAX_OUTPUT_LEN",")",">",">",">","outputs","=","infer_engine.generate","(","input_ids",",","*","*","generate_kwargs",")"]
docstr: Generate output tokens by setting BatchInferState as an attribute to the model and calling model.generate
  
  Args:
      inputs: should be one of the following types
          1. BatchEncoding or dict (e.g. tokenizer batch_encode)
          2. list of input token ids (e.g. appended result of tokenizer encode)
          3. torch.Tensor (e.g. tokenizer encode with return_tensors='pt')
docstr_tok: ["Generate","output","tokens","by","setting","BatchInferState","as","an","attribute","to","the","model","and","calling","model.generate","Args",":","inputs",":","should","be","one","of","the","following","types","1",".","BatchEncoding","or","dict","(","e.g",".","tokenizer","batch_encode",")","2.","list","of","input","token","ids","(","e.g",".","appended","result","of","tokenizer","encode",")","3.","torch.Tensor","(","e.g",".","tokenizer","encode","with","return_tensors='pt","'",")"]
code: def _generate_by_set_infer_state(
      self, input_tokens, **generate_kwargs
  ) -> torch.Tensor:
      """
      Generate output tokens by setting BatchInferState as an attribute to the model and calling model.generate
  
      Args:
          inputs: should be one of the following types
              1. BatchEncoding or dict (e.g. tokenizer batch_encode)
              2. list of input token ids (e.g. appended result of tokenizer encode)
              3. torch.Tensor (e.g. tokenizer encode with return_tensors='pt')
      """
  
      # for testing, always use sharded model
      assert self.model is not None, "sharded model does not exist"
  
      batch_infer_state = self.prepare_batch_state(input_tokens)
      assert (
          batch_infer_state.max_len_in_batch <= self.max_input_len
      ), "max length in batch exceeds limit"
  
      # set BatchInferState for the current batch as attr to model
      # NOTE this is not a preferable way to pass BatchInferState during inference
      #   we might want to rewrite generate function (e.g. _generate_by_pass_infer_state)
      #   and pass BatchInferState via model forward
      model = self.model
      if isinstance(model, LlamaForCausalLM):
          model = self.model.model
      elif isinstance(model, BloomForCausalLM):
          model = self.model.transformer
      setattr(model, "infer_state", batch_infer_state)
  
      outputs = self.model.generate(
          **input_tokens, **generate_kwargs, early_stopping=False
      )
  
      # NOTE In future development, we're going to let the scheduler to handle the cache,
      #      instead of freeing space explicitly at the end of generation
      self.cache_manager.free_all()
  
      return outputs
  
imports: from typing import Any, Callable, List, Optional, Union
  import torch
  import torch.nn
  from transformers import BloomForCausalLM, LlamaForCausalLM
  from transformers.generation import GenerationConfig
  from transformers.generation.stopping_criteria import StoppingCriteriaList
  from transformers.tokenization_utils_base import BatchEncoding
  from colossalai.shardformer import ShardConfig, ShardFormer
  from colossalai.shardformer.policies.auto_policy import get_autopolicy
  from .batch_infer_state import BatchInferState
  from .kvcache_manager import MemoryManager
label_desc: Use image node_id 10 for calling the TPInferEngine obj's underlying member method code with example usage: obj._generate_by_set_infer_state(input_tokens) and returns: outputs