code: def _replace(self, child, name, conv_linear_layer):
      if getattr(child, "replaced", False) == True:
          return
      weight_shape = child.weight.shape
      mp_replace = ReplaceWithTensorSlicing(mp_group=self.mp_group)
      if name in self.all_reduce_linears:
          # if conv_linear_layer [weight_shape[1], weight_shape[0] // mp_size]
          # else [weight_shape[0], weight_shape[1] // mp_size]
  
          if self.conv_linear_layer:
              child.weight.data = child.weight.data.transpose(
                  -1, -2
              ).contiguous()
          data = child.weight.data.split(
              get_shard_size_list(
                  weight_shape[0]
                  if self.conv_linear_layer
                  else weight_shape[1],
                  self.mp_size,
                  name,
              ),
              dim=1,
          )
          data_dc = move(
              data[mp_replace.gpu_index],
              get_accelerator().current_device_name(),
          ).detach()
          del data
  
          setattr(child, "replaced", True)
          if name == "lm_head" or name == "embed_out":
              return LmHeadLinearAllreduce(
                  torch.nn.parameter.Parameter(
                      data_dc, requires_grad=False
                  ),
                  dist.get_rank(),
                  dist.get_world_size(),
                  child.bias
                  if child.bias is None
                  else torch.nn.parameter.Parameter(
                      move(
                          child.bias,
                          get_accelerator().current_device_name(),
                      )
                  ),
                  self.mp_group,
              )
          return LinearAllreduce(
              torch.nn.parameter.Parameter(
                  data_dc, requires_grad=False
              ),
              child.bias
              if child.bias is None
              else torch.nn.parameter.Parameter(
                  move(
                      child.bias,
                      get_accelerator().current_device_name(),
                  )
              ),
              self.mp_group,
          )
      else:
          # if conv_linear_layer [weight_shape[1], weight_shape[0] // mp_size]
          # else [weight_shape[0] // mp_size, weight_shape[1]]
          if self.conv_linear_layer:
              child.weight.data = child.weight.data.transpose(
                  -1, -2
              ).contiguous()
  
          if require_tp_fused_qkvw(name, self.mp_size):
              # Check and handle fused qkv for TP
              # The copy is a regular copy, The shape of dst and src is the same
              data_dc = move(
                  prepare_tp_fused_qkvw(
                      self.module,
                      child.weight.data,
                      self.mp_size,
                      mp_replace.gpu_index,
                  ),
                  get_accelerator().current_device_name(),
              )
  
              bias_data_dc = (
                  None
                  if child.bias is None
                  else move(
                      prepare_tp_fused_qkvw(
                          self.module,
                          child.bias.data,
                          self.mp_size,
                          mp_replace.gpu_index,
                      ),
                      get_accelerator().current_device_name(),
                  )
              )
          else:
              data = child.weight.data.split(
                  get_shard_size_list(
                      weight_shape[0], self.mp_size, name
                  ),
                  dim=1 if self.conv_linear_layer else 0,
              )
              data_dc = move(
                  data[mp_replace.gpu_index],
                  get_accelerator().current_device_name(),
              ).detach()
              del data
  
              if child.bias is not None:
                  bias_data = child.bias.data.split(
                      get_shard_size_list(
                          weight_shape[1]
                          if self.conv_linear_layer
                          else weight_shape[0],
                          self.mp_size,
                          name,
                      ),
                      dim=0,
                  )
                  bias_data = move(
                      bias_data[mp_replace.gpu_index],
                      get_accelerator().current_device_name(),
                  )
                  bias_data_dc = torch.nn.parameter.Parameter(
                      bias_data, requires_grad=False
                  )
                  del bias_data
              else:
                  bias_data_dc = None
  
          setattr(child, "replaced", True)
          return LinearLayer(
              weight=torch.nn.parameter.Parameter(
                  data_dc, requires_grad=False
              ),
              bias=bias_data_dc,
          )
  
imports: import re
  from torch import nn
  from .replace_policy import replace_policies
  from typing import Optional
  import torch
  from deepspeed import comm
  from .layers import LinearAllreduce, LinearLayer, LmHeadLinearAllreduce
  from deepspeed.accelerator import get_accelerator
  from .fusedqkv_utils import require_tp_fused_qkvw, prepare_tp_fused_qkvw
  from deepspeed.module_inject.tp_shard import get_shard_size, get_shard_size_list
label_desc: Use image node_id 10 for calling the AutoTP obj's underlying member method code with example usage: obj._replace(child, name, conv_linear_layer) and returns: LinearAllreduce, LinearLayer, LmHeadLinearAllreduce