code: def _rewrite_spec_if_needed(
      spec: shard_spec.ShardingSpec, tensor: torch.Tensor, rank: int
  ) -> shard_spec.ShardingSpec:
      """
      Rewrite ``spec`` to match the device of ``tensor``.
  
      FSDP.sharded_optim_state_dict sneakly ships optimizer state to CPU so if the original ShardingSpec
      produces CUDA metadata, ST construction bombs.
      """
      if not isinstance(spec, ChunkShardingSpec):
          return spec
  
      # let's see if we need
      rewrite = False
      for p in spec.placements:
          p = cast(_remote_device, p)
          if p.rank() == rank and p.device() != tensor.device:
              rewrite = True
              break
      if rewrite:
          spec = copy.deepcopy(spec)
          for i, placement in enumerate(spec.placements):
              placement = cast(_remote_device, placement)
              if (
                  placement.rank() == rank
                  and placement.device() != tensor.device
              ):
                  spec.placements[i] = _remote_device(
                      f"rank:{rank}/{tensor.device}"
                  )
  
      return spec
  
imports: import copy
  from typing import Any, cast, List, Optional, Tuple
  import torch
  import torch.distributed
  import torch.distributed._shard.sharding_spec
  import torch.distributed.distributed_c10d
  from torch.distributed._shard.sharded_tensor import Shard, ShardedTensor, ShardedTensorMetadata, TensorProperties
  from torch.distributed._shard.sharding_spec import ShardMetadata
  from torch.distributed._shard.sharding_spec.chunk_sharding_spec import ChunkShardingSpec
  from torch.distributed._tensor import DeviceMesh, DTensor, Replicate, Shard
  from torch.distributed.device_mesh import _mesh_resources
  from torch.distributed.fsdp._common_utils import _set_fsdp_flattened
  from torch.distributed.fsdp._fsdp_extensions import FSDPExtensions
  from torch.distributed.fsdp._shard_utils import _create_chunk_sharded_tensor
  from torch.distributed.remote_device import _remote_device
  from torch.distributed.tensor.parallel._data_parallel_utils import _flatten_tensor, _unflatten_tensor
label_desc: Use image node_id 7 for calling a global function with example usage: _rewrite_spec_if_needed(spec, tensor, rank) and returns: spec, spec