import torch

from accelerate import Accelerator, DistributedType







class __snake_case							:




   def __enter__(       self     :							str ):



     """simple docstring"""






     if self.enabled:
       SCREAMING_SNAKE_CASE__     				= self.model.no_sync()
       self.model_sync_obj.__enter__()

     return self




   def __exit__(       self     :							Optional[Any]			,      _lowercase     :							List[Any]			,      _lowercase     :							List[str]			,      _lowercase     :							Tuple ):



     """simple docstring"""






     if self.enabled:
       # Average all models on exit
       self._sync_and_avg_model_params()
       self.model_sync_obj.__exit__(_lowercase			,      _lowercase			,      _lowercase )




   def __init__(       self     :							Tuple			,      _lowercase     :							Accelerator			,      _lowercase     :							torch.nn.Module			,      _lowercase     :							int			,      _lowercase     :							bool = True ):



     """simple docstring"""






     if accelerator.distributed_type not in [
         DistributedType.NO,
         DistributedType.MULTI_CPU,
         DistributedType.MULTI_GPU,
     ]:
       raise NotImplementedError("""LocalSGD is supported only for CPUs and GPUs (no DeepSpeed or MegatronLM)""" )
     SCREAMING_SNAKE_CASE__     				= enabled and accelerator.distributed_type != DistributedType.NO
     SCREAMING_SNAKE_CASE__     				= 0
     if self.enabled:
       SCREAMING_SNAKE_CASE__     				= accelerator
       SCREAMING_SNAKE_CASE__     				= model
       SCREAMING_SNAKE_CASE__     				= local_sgd_steps




   def       __a		(       self     :							int ):



     """simple docstring"""






     self.num_steps += 1
     if not self.enabled:
       return

     if self.num_steps % self.local_sgd_steps == 0:
       self._sync_and_avg_model_params()




   def       __a		(       self     :							Tuple ):



     """simple docstring"""






     self.accelerator.wait_for_everyone()
     with self.accelerator.autocast():
       for param in self.model.parameters():
         SCREAMING_SNAKE_CASE__     				= self.accelerator.reduce(param.data			,      reduction="""mean""" )

from typing import TYPE_CHECKING

from ...utils import OptionalDependencyNotAvailable, _LazyModule, is_torch_available, is_vision_available


__lowerCamelCase     :					Optional[int]				  =						{'''configuration_vilt''': ['''VILT_PRETRAINED_CONFIG_ARCHIVE_MAP''', '''ViltConfig''']}

try:
   if not is_vision_available():
      raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
   pass
else:
   __lowerCamelCase     :					Dict				  =						['''ViltFeatureExtractor''']
   __lowerCamelCase     :					Dict				  =						['''ViltImageProcessor''']
   __lowerCamelCase     :					str				  =						['''ViltProcessor''']

try:
   if not is_torch_available():
      raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
   pass
else:
   __lowerCamelCase     :					Any				  =						[
       '''VILT_PRETRAINED_MODEL_ARCHIVE_LIST''',
       '''ViltForImageAndTextRetrieval''',
       '''ViltForImagesAndTextClassification''',
       '''ViltForTokenClassification''',
       '''ViltForMaskedLM''',
       '''ViltForQuestionAnswering''',
       '''ViltLayer''',
       '''ViltModel''',
       '''ViltPreTrainedModel''',
   ]


if TYPE_CHECKING:
   from .configuration_vilt import VILT_PRETRAINED_CONFIG_ARCHIVE_MAP, ViltConfig

   try:
      if not is_vision_available():
         raise OptionalDependencyNotAvailable()
   except OptionalDependencyNotAvailable:
      pass
   else:
      from .feature_extraction_vilt import ViltFeatureExtractor
      from .image_processing_vilt import ViltImageProcessor
      from .processing_vilt import ViltProcessor

   try:
      if not is_torch_available():
         raise OptionalDependencyNotAvailable()
   except OptionalDependencyNotAvailable:
      pass
   else:
      from .modeling_vilt import (
          VILT_PRETRAINED_MODEL_ARCHIVE_LIST,
          ViltForImageAndTextRetrieval,
          ViltForImagesAndTextClassification,
          ViltForMaskedLM,
          ViltForQuestionAnswering,
          ViltForTokenClassification,
          ViltLayer,
          ViltModel,
          ViltPreTrainedModel,
      )


else:
   import sys

   __lowerCamelCase     :					List[str]				  =						_LazyModule(__name__, globals()['''__file__'''], _import_structure)
