code: def _unflatten_tensor(
      tensor, spec, *, device_handle=None, compute_stream=None
  ):
      # unflatten would mainly be called everytime FSDP allgather parameters.
      result = DTensor.from_local(
          tensor,
          spec.mesh,
          spec.placements,
          run_check=False,
          shape=spec.shape,
          stride=spec.stride,
      )
      if tensor.requires_grad:
          # only register the hook if the tensor requires grad
          tensor.register_hook(
              partial(
                  sync_grad_hook,
                  device_handle=device_handle,
                  compute_stream=compute_stream,
              )
          )
      return result
  
file_path: pytorch/torch/distributed/tensor/parallel/_data_parallel_utils.py
imports: from functools import partial
  from typing import no_type_check, Optional, Tuple
  import torch
  from torch.distributed._functional_collectives import AsyncCollectiveTensor
  from torch.distributed._tensor import DTensor
  from torch.distributed._tensor.placement_types import DTensorSpec
label_desc: Use image node_id 3 for calling a global function with example usage: _unflatten_tensor(tensor, spec) and returns: result