code: def with_xla(func: Callable) -> Callable:
      assert func is not None
  
      @wraps(func)  # pyre-ignore[6]
      def wrapper(
          self, *args: Tuple[object], **kwargs: Dict[str, Any]  # type: ignore[misc]
      ) -> None:
          if TORCH_XLA_INITIALIZED:
              # TODO(yeounoh) replace this with xr.use_spmd() when we deprecate the flag.
              os.environ["XLA_USE_SPMD"] = "1"
              return func(self, *args, **kwargs)  # type: ignore[misc]
          else:
              raise ImportError(
                  "torch.distributed._tensor._xla API requires torch_xla package installation."
              )
  
      return wrapper
  
file_path: pytorch/torch/distributed/_tensor/_xla.py
imports: import logging
  import os
  from functools import wraps
  from typing import Any, Callable, Dict, Optional, Sequence, Tuple, Union
  import torch
  import torch.nn
  from torch.distributed._tensor.placement_types import Placement, Replicate
  from torch.distributed.device_mesh import DeviceMesh
label_desc: Use image node_id 1 for calling a global function with example usage: with_xla(func) and returns: wrapper, func