code: def weight_split(
      weight: torch.nn.parameter.Parameter, dim: int, col_normal: bool
  ) -> torch.nn.parameter.Parameter:
      """weight_split
      split a nn.Parameter
  
      Args:
          weight (torch.nn.parameter.Parameter): a torch Parameter instance
          dim (int): the dimension to be sharded along with
          col_normal(bool): col shard with gather or not
      Returns:
          _type_: _description_
      """
      if col_normal:
          setattr(weight, "fx_attr", (dim, "SHARD", "TP", "col_normal"))
      else:
          setattr(
              weight,
              "fx_attr",
              (dim, "SHARD", "TP", "col_needs_many_outputs"),
          )
      return weight
  
file_path: colossalai/colossalai/fx/passes/shard_1d_pass.py
imports: import operator
  import torch
  import torch.nn
  from colossalai.legacy.tensor import ProcessGroup
  from colossalai.legacy.tensor.compute_spec import ComputePattern, ComputeSpec
  from colossalai.legacy.tensor.distspec import ShardSpec
label_desc: Use image node_id 1 for calling a global function with example usage: weight_split(weight, dim, col_normal) and returns: weight