code: def _apply_tp_linear(
      linear: nn.Linear, style: str, weight_splits: List[int] = []
  ) -> None:
      rank = _get_rank()
      world_size = _get_world_size()
  
      # Linear's weight matrix is transposed, and is of shape
      # (linear.out_features, linear.in_features)
      dim_lookup = {
          "colwise": (0, "out_features"),
          "rowwise": (1, "in_features"),
      }
      assert style in dim_lookup
      shard_dim, size_attr = dim_lookup[style]
  
      # ensure we can shard evenly
      assert getattr(linear, size_attr) % world_size == 0
  
      def shard(x, dim):
          assert x.size(dim=dim) % world_size == 0
          return torch.tensor_split(x, world_size, dim=dim)[rank]
  
      def shard_qkv(qkv, dim, weight_splits):
          q, k, v = qkv.split(weight_splits, dim=dim)
          q = shard(q, dim)
          k = shard(k, dim)
          v = shard(v, dim)
          return torch.cat((q, k, v), dim=dim)
  
      # shard
      if weight_splits:
          # attention
          assert len(weight_splits) == 3
  
          if isinstance(linear, WeightOnlyInt4Linear):
              sharded_weight = shard_qkv(
                  linear.weight,
                  shard_dim,
                  [i // 8 for i in weight_splits],
              )
              linear.scales_and_zeros = shard_qkv(
                  linear.scales_and_zeros, 1 - shard_dim, weight_splits
              )
          else:
              sharded_weight = shard_qkv(
                  linear.weight, shard_dim, weight_splits
              )
          if hasattr(linear, "scales") and style == "colwise":
              linear.scales = shard_qkv(linear.scales, 0, weight_splits)
      else:
          sharded_weight = shard(linear.weight, shard_dim)
          if isinstance(linear, WeightOnlyInt4Linear):
              linear.scales_and_zeros = shard(
                  linear.scales_and_zeros, 1 - shard_dim
              )
              if style == "rowwise":
                  assert (
                      linear.scales_and_zeros.shape[0] * 32
                      == sharded_weight.shape[1]
                      * sharded_weight.shape[2]
                      * sharded_weight.shape[3]
                  )
                  assert (
                      linear.scales_and_zeros.shape[1]
                      == sharded_weight.shape[0] * 8
                  )
          if hasattr(linear, "scales") and style == "colwise":
              linear.scales = shard(linear.scales, 0)
  
      # local_break()
      linear.weight = nn.Parameter(sharded_weight, requires_grad=False)
      setattr(
          linear, size_attr, getattr(linear, size_attr) // world_size
      )
  
file_path: gpt-fast/tp.py
imports: import os
  from typing import List, Optional
  import torch
  import torch.distributed
  from torch import nn
  from torch.distributed import _functional_collectives
  from model import Attention, FeedForward, Transformer
  from quantize import WeightOnlyInt4Linear
label_desc: Use image node_id 6 for calling a global function with example usage: _apply_tp_linear(linear, style, weight_splits) and returns: torch, torch