






from typing import TYPE_CHECKING

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


UpperCAmelCase_ :						str      =      {'''configuration_deit''': ['''DEIT_PRETRAINED_CONFIG_ARCHIVE_MAP''', '''DeiTConfig''', '''DeiTOnnxConfig''']}

try:
	if not is_vision_available():
		raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
	pass
else:
	UpperCAmelCase_ :						Dict      =      ['''DeiTFeatureExtractor''']
	UpperCAmelCase_ :						int      =      ['''DeiTImageProcessor''']

try:
	if not is_torch_available():
		raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
	pass
else:
	UpperCAmelCase_ :						Optional[int]      =      [
	    '''DEIT_PRETRAINED_MODEL_ARCHIVE_LIST''',
	    '''DeiTForImageClassification''',
	    '''DeiTForImageClassificationWithTeacher''',
	    '''DeiTForMaskedImageModeling''',
	    '''DeiTModel''',
	    '''DeiTPreTrainedModel''',
	]

try:
	if not is_tf_available():
		raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
	pass
else:
	UpperCAmelCase_ :						Any      =      [
	    '''TF_DEIT_PRETRAINED_MODEL_ARCHIVE_LIST''',
	    '''TFDeiTForImageClassification''',
	    '''TFDeiTForImageClassificationWithTeacher''',
	    '''TFDeiTForMaskedImageModeling''',
	    '''TFDeiTModel''',
	    '''TFDeiTPreTrainedModel''',
	]


if TYPE_CHECKING:
	from .configuration_deit import DEIT_PRETRAINED_CONFIG_ARCHIVE_MAP, DeiTConfig, DeiTOnnxConfig

	try:
		if not is_vision_available():
			raise OptionalDependencyNotAvailable()
	except OptionalDependencyNotAvailable:
		pass
	else:
		from .feature_extraction_deit import DeiTFeatureExtractor
		from .image_processing_deit import DeiTImageProcessor

	try:
		if not is_torch_available():
			raise OptionalDependencyNotAvailable()
	except OptionalDependencyNotAvailable:
		pass
	else:
		from .modeling_deit import (
		    DEIT_PRETRAINED_MODEL_ARCHIVE_LIST,
		    DeiTForImageClassification,
		    DeiTForImageClassificationWithTeacher,
		    DeiTForMaskedImageModeling,
		    DeiTModel,
		    DeiTPreTrainedModel,
		)

	try:
		if not is_tf_available():
			raise OptionalDependencyNotAvailable()
	except OptionalDependencyNotAvailable:
		pass
	else:
		from .modeling_tf_deit import (
		    TF_DEIT_PRETRAINED_MODEL_ARCHIVE_LIST,
		    TFDeiTForImageClassification,
		    TFDeiTForImageClassificationWithTeacher,
		    TFDeiTForMaskedImageModeling,
		    TFDeiTModel,
		    TFDeiTPreTrainedModel,
		)


else:
	import sys

	UpperCAmelCase_ :						Any      =      _LazyModule(__name__, globals()['''__file__'''], _import_structure, module_spec=__spec__)





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

import torch_xla.distributed.xla_multiprocessing as xmp


def 			UpperCamelCase(  ):
		lowerCAmelCase_			:     List[str]									= 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 			UpperCamelCase(  ):
		lowerCAmelCase_			:     str									= parse_args()

		# Import training_script as a module.
		lowerCAmelCase_			:     str									= Path(args.training_script )
		sys.path.append(str(script_fpath.parent.resolve() ) )
		lowerCAmelCase_			:     Tuple									= script_fpath.stem
		lowerCAmelCase_			:     Union[str, Any]									= importlib.import_module(__UpperCamelCase )

		# Patch sys.argv
		lowerCAmelCase_			:     Optional[int]									= [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()





