code: def search_tp_partition_dim(
      current_shape: torch.Size,
      original_shape: torch.Size,
      tp_size: int,
  ) -> Optional[int]:
      """
      Given the current shape of parameter and the shape of parameter before sharding,
      return the dimension along which the parameter is sharded when using tensor parallel.
      If tensor parallel is not used, return None.
  
      Args:
          current_shape (torch.Size): The current shape of parameter after sharding.
          original_shape (torch.Size): The shape of parameter before sharding.
          tp_size (int): The size of tp group.
  
      Returns:
          Optional[int]: The dimension along which parameter is partitioned.
      """
      partition_dim = None
      for dim, length in enumerate(original_shape):
          if length > current_shape[dim]:
              partition_dim = dim
              break
      if partition_dim is not None:
          assert (
              original_shape[partition_dim]
              == tp_size * current_shape[partition_dim]
          ), f"The parameter isn't evenly distributed among tensor parallel group: \
                  shape before sharding {original_shape}, shape after sharding {current_shape}"
  
      return partition_dim
  
imports: import os
  import re
  from collections import abc
  from collections import defaultdict
  from itertools import chain
  from pathlib import Path
  from typing import Iterator, List, Mapping, Optional, OrderedDict, Tuple
  import torch
  import torch.nn
  from packaging.version import Version
  from torch.optim import Optimizer
  from torch.utils._pytree import tree_map
  from colossalai.tensor.d_tensor import is_customized_distributed_tensor, is_distributed_tensor, to_global, to_global_for_customized_distributed_tensor
label_desc: Use image node_id 5 for calling a global function with example usage: search_tp_partition_dim(current_shape, original_shape, tp_size) and returns: partition_dim