

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

import torch_xla.distributed.xla_multiprocessing as xmp




def        __UpperCamelCase     (      )				->List[Any]:



							"""simple docstring"""
							lowerCamelCase_								=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=_A      ,  default=1      ,  help="""Number of TPU cores to use (1 or 8)."""		)

							# positional
							parser.add_argument(
							    """training_script"""      ,  type=_A      ,  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=_A		)

							return parser.parse_args()




def        __UpperCamelCase     (      )				->List[str]:



							"""simple docstring"""
							lowerCamelCase_								=parse_args()

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

							# Patch sys.argv
							lowerCamelCase_								=[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_torch_available,
    is_vision_available,
)


__A	:   Optional[int]	=					{
    'configuration_mobilevitv2': [
        'MOBILEVITV2_PRETRAINED_CONFIG_ARCHIVE_MAP',
        'MobileViTV2Config',
        'MobileViTV2OnnxConfig',
    ],
}


try:
	if not is_torch_available():
		raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
	pass
else:
	__A	:   Tuple	=					[
	    'MOBILEVITV2_PRETRAINED_MODEL_ARCHIVE_LIST',
	    'MobileViTV2ForImageClassification',
	    'MobileViTV2ForSemanticSegmentation',
	    'MobileViTV2Model',
	    'MobileViTV2PreTrainedModel',
	]

if TYPE_CHECKING:
	from .configuration_mobilevitva import (
	    MOBILEVITV2_PRETRAINED_CONFIG_ARCHIVE_MAP,
	    MobileViTVaConfig,
	    MobileViTVaOnnxConfig,
	)

	try:
		if not is_torch_available():
			raise OptionalDependencyNotAvailable()
	except OptionalDependencyNotAvailable:
		pass
	else:
		from .modeling_mobilevitva import (
		    MOBILEVITV2_PRETRAINED_MODEL_ARCHIVE_LIST,
		    MobileViTVaForImageClassification,
		    MobileViTVaForSemanticSegmentation,
		    MobileViTVaModel,
		    MobileViTVaPreTrainedModel,
		)

else:
	import sys

	__A	:   List[str]	=					_LazyModule(__name__, globals()['__file__'], _import_structure, module_spec=__spec__)
