



import importlib
import sys
from argparse import REMAINDER, ArgumentParser
from pathlib import Path

import torch_xla.distributed.xla_multiprocessing as xmp





def   __snake_case			(    ):
    __a		 =       ArgumentParser(
        description=(
            '''PyTorch TPU distributed training launch '''
            '''helper utility that will spawn up '''
            '''multiple distributed processes'''
        )       )

    # Optional arguments for the launch helper
    parser.add_argument('''--num_cores'''     ,				type=_UpperCAmelCase     ,				default=1     ,				help='''Number of TPU cores to use (1 or 8).'''       )

    # positional
    parser.add_argument(
        '''training_script'''     ,				type=_UpperCAmelCase     ,				help=(
            '''The full path to the single TPU training '''
            '''program/script to be launched in parallel, '''
            '''followed by all the arguments for the '''
            '''training script'''
        )     ,				)

    # rest from the training program
    parser.add_argument('''training_script_args'''     ,				nargs=_UpperCAmelCase       )

    return parser.parse_args()





def   __snake_case			(    ):
    __a		 =       parse_args()

    # Import training_script as a module.
    __a		 =       Path(args.training_script       )
    sys.path.append(str(script_fpath.parent.resolve()       )       )
    __a		 =       script_fpath.stem
    __a		 =       importlib.import_module(_UpperCAmelCase       )

    # Patch sys.argv
    __a		 =       [args.training_script] + args.training_script_args + ['''--tpu_num_cores''', str(args.num_cores       )]
    xmp.spawn(mod._mp_fn     ,				args=()     ,				nprocs=args.num_cores       )


if __name__ == "__main__":
 main()










from typing import TYPE_CHECKING

from ...utils import (
    OptionalDependencyNotAvailable,
    _LazyModule,
    is_flax_available,
    is_tf_available,
    is_torch_available,
    is_vision_available,
)


__snake_case		:Union[str, Any]   	=  {'''configuration_vit''': ['''VIT_PRETRAINED_CONFIG_ARCHIVE_MAP''', '''ViTConfig''', '''ViTOnnxConfig''']}

try:
 if not is_vision_available():
  raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
 pass
else:
 __snake_case		:List[str]   	=  ['''ViTFeatureExtractor''']
 __snake_case		:Optional[Any]   	=  ['''ViTImageProcessor''']

try:
 if not is_torch_available():
  raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
 pass
else:
 __snake_case		:str   	=  [
     '''VIT_PRETRAINED_MODEL_ARCHIVE_LIST''',
     '''ViTForImageClassification''',
     '''ViTForMaskedImageModeling''',
     '''ViTModel''',
     '''ViTPreTrainedModel''',
 ]

try:
 if not is_tf_available():
  raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
 pass
else:
 __snake_case		:Tuple   	=  [
     '''TFViTForImageClassification''',
     '''TFViTModel''',
     '''TFViTPreTrainedModel''',
 ]

try:
 if not is_flax_available():
  raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
 pass
else:
 __snake_case		:Tuple   	=  [
     '''FlaxViTForImageClassification''',
     '''FlaxViTModel''',
     '''FlaxViTPreTrainedModel''',
 ]

if TYPE_CHECKING:
 from .configuration_vit import VIT_PRETRAINED_CONFIG_ARCHIVE_MAP, ViTConfig, ViTOnnxConfig

 try:
  if not is_vision_available():
   raise OptionalDependencyNotAvailable()
 except OptionalDependencyNotAvailable:
  pass
 else:
  from .feature_extraction_vit import ViTFeatureExtractor
  from .image_processing_vit import ViTImageProcessor

 try:
  if not is_torch_available():
   raise OptionalDependencyNotAvailable()
 except OptionalDependencyNotAvailable:
  pass
 else:
  from .modeling_vit import (
      VIT_PRETRAINED_MODEL_ARCHIVE_LIST,
      ViTForImageClassification,
      ViTForMaskedImageModeling,
      ViTModel,
      ViTPreTrainedModel,
  )

 try:
  if not is_tf_available():
   raise OptionalDependencyNotAvailable()
 except OptionalDependencyNotAvailable:
  pass
 else:
  from .modeling_tf_vit import TFViTForImageClassification, TFViTModel, TFViTPreTrainedModel

 try:
  if not is_flax_available():
   raise OptionalDependencyNotAvailable()
 except OptionalDependencyNotAvailable:
  pass
 else:
  from .modeling_flax_vit import FlaxViTForImageClassification, FlaxViTModel, FlaxViTPreTrainedModel


else:
 import sys

 __snake_case		:Optional[int]   	=  _LazyModule(__name__, globals()['''__file__'''], _import_structure, module_spec=__spec__)





