code: def _is_tp_sharded(fqn: str) -> bool:
      """
      Returns whether a tensor given by the fqn is tensor parallel sharded.
      NOTE: this is currently done by inspection of the MF model and is quite
      brittle and would need to be updated if the MF sharding changes.
      """
      return (
          "attention" in fqn
          or "feed_forward" in fqn
          or "output" in fqn
          or "tok_embeddings" in fqn
      )
  
file_path: serve/examples/large_models/tp_llama/checkpoint_converter.py
imports: import logging
  from dataclasses import dataclass
  from typing import Dict, List, Union
  import torch
  import torch.distributed
  from torch import nn, Tensor
  from torch.distributed._shard.sharded_tensor.api import ShardedTensor
  from torch.distributed._tensor import DeviceMesh, DTensor
  from torch.distributed.fsdp._fsdp_extensions import _ext_chunk_dtensor, _ext_chunk_tensor
label_desc: Use image node_id 3 for calling a global function with example usage: _is_tp_sharded(fqn) and returns: unknown