




import collections
import inspect
import unittest

from transformers import SwinConfig
from transformers.testing_utils import require_torch, require_vision, slow, torch_device
from transformers.utils import cached_property, is_torch_available, is_vision_available

from ...test_backbone_common import BackboneTesterMixin
from ...test_configuration_common import ConfigTester
from ...test_modeling_common import ModelTesterMixin, _config_zero_init, floats_tensor, ids_tensor
from ...test_pipeline_mixin import PipelineTesterMixin


if is_torch_available():
  import torch
  from torch import nn

  from transformers import SwinBackbone, SwinForImageClassification, SwinForMaskedImageModeling, SwinModel
  from transformers.models.swin.modeling_swin import SWIN_PRETRAINED_MODEL_ARCHIVE_LIST


if is_vision_available():
  from PIL import Image

  from transformers import AutoImageProcessor
class       _UpperCamelCase   :
      def __init__(      self :str			,   lowerCamelCase :str			,   lowerCamelCase :Any=13			,   lowerCamelCase :Dict=32			,   lowerCamelCase :Dict=2			,   lowerCamelCase :Optional[Any]=3			,   lowerCamelCase :Tuple=16			,   lowerCamelCase :Union[str, Any]=[1, 2, 1]			,   lowerCamelCase :Union[str, Any]=[2, 2, 4]			,   lowerCamelCase :Any=2			,   lowerCamelCase :Tuple=2.0			,   lowerCamelCase :Optional[Any]=True			,   lowerCamelCase :Optional[int]=0.0			,   lowerCamelCase :Any=0.0			,   lowerCamelCase :Optional[Any]=0.1			,   lowerCamelCase :Optional[Any]="gelu"			,   lowerCamelCase :Optional[Any]=False			,   lowerCamelCase :Dict=True			,   lowerCamelCase :Union[str, Any]=0.02			,   lowerCamelCase :Union[str, Any]=1e-5			,   lowerCamelCase :Tuple=True			,   lowerCamelCase :str=None			,   lowerCamelCase :Any=True			,   lowerCamelCase :List[str]=10			,   lowerCamelCase :Tuple=8			,   lowerCamelCase :Any=["stage1", "stage2"]			,   lowerCamelCase :Optional[int]=[1, 2]			,   )  ->      Tuple:
       UpperCAmelCase__							=    parent
       UpperCAmelCase__							=    batch_size
       UpperCAmelCase__							=    image_size
       UpperCAmelCase__							=    patch_size
       UpperCAmelCase__							=    num_channels
       UpperCAmelCase__							=    embed_dim
       UpperCAmelCase__							=    depths
       UpperCAmelCase__							=    num_heads
       UpperCAmelCase__							=    window_size
       UpperCAmelCase__							=    mlp_ratio
       UpperCAmelCase__							=    qkv_bias
       UpperCAmelCase__							=    hidden_dropout_prob
       UpperCAmelCase__							=    attention_probs_dropout_prob
       UpperCAmelCase__							=    drop_path_rate
       UpperCAmelCase__							=    hidden_act
       UpperCAmelCase__							=    use_absolute_embeddings
       UpperCAmelCase__							=    patch_norm
       UpperCAmelCase__							=    layer_norm_eps
       UpperCAmelCase__							=    initializer_range
       UpperCAmelCase__							=    is_training
       UpperCAmelCase__							=    scope
       UpperCAmelCase__							=    use_labels
       UpperCAmelCase__							=    type_sequence_label_size
       UpperCAmelCase__							=    encoder_stride
       UpperCAmelCase__							=    out_features
       UpperCAmelCase__							=    out_indices
      def 						UpperCAmelCase_    (      self :Optional[int]   )  ->      Dict:
       UpperCAmelCase__							=    floats_tensor([self.batch_size, self.num_channels, self.image_size, self.image_size]   )

       UpperCAmelCase__							=    None
       if self.use_labels:
        UpperCAmelCase__							=    ids_tensor([self.batch_size]			,   self.type_sequence_label_size   )

       UpperCAmelCase__							=    self.get_config()

       return config, pixel_values, labels
      def 						UpperCAmelCase_    (      self :List[str]   )  ->      str:
       return SwinConfig(
           image_size=self.image_size			,   patch_size=self.patch_size			,   num_channels=self.num_channels			,   embed_dim=self.embed_dim			,   depths=self.depths			,   num_heads=self.num_heads			,   window_size=self.window_size			,   mlp_ratio=self.mlp_ratio			,   qkv_bias=self.qkv_bias			,   hidden_dropout_prob=self.hidden_dropout_prob			,   attention_probs_dropout_prob=self.attention_probs_dropout_prob			,   drop_path_rate=self.drop_path_rate			,   hidden_act=self.hidden_act			,   use_absolute_embeddings=self.use_absolute_embeddings			,   path_norm=self.patch_norm			,   layer_norm_eps=self.layer_norm_eps			,   initializer_range=self.initializer_range			,   encoder_stride=self.encoder_stride			,   out_features=self.out_features			,   out_indices=self.out_indices			,   )
      def 						UpperCAmelCase_    (      self :int			,   lowerCamelCase :List[str]			,   lowerCamelCase :str			,   lowerCamelCase :List[Any]   )  ->      Tuple:
       UpperCAmelCase__							=    SwinModel(config=lowerCamelCase   )
       model.to(lowerCamelCase   )
       model.eval()
       UpperCAmelCase__							=    model(lowerCamelCase   )

       UpperCAmelCase__							=    ((config.image_size // config.patch_size) ** 2) // (4 ** (len(config.depths   ) - 1))
       UpperCAmelCase__							=    int(config.embed_dim * 2 ** (len(config.depths   ) - 1)   )

       self.parent.assertEqual(result.last_hidden_state.shape			,   (self.batch_size, expected_seq_len, expected_dim)   )
      def 						UpperCAmelCase_    (      self :Dict			,   lowerCamelCase :int			,   lowerCamelCase :List[Any]			,   lowerCamelCase :str   )  ->      Any:
       UpperCAmelCase__							=    SwinBackbone(config=lowerCamelCase   )
       model.to(lowerCamelCase   )
       model.eval()
       UpperCAmelCase__							=    model(lowerCamelCase   )

       # verify hidden states
       self.parent.assertEqual(len(result.feature_maps   )			,   len(config.out_features   )   )
       self.parent.assertListEqual(list(result.feature_maps[0].shape   )			,   [self.batch_size, model.channels[0], 16, 16]   )

       # verify channels
       self.parent.assertEqual(len(model.channels   )			,   len(config.out_features   )   )

       # verify backbone works with out_features=None
       UpperCAmelCase__							=    None
       UpperCAmelCase__							=    SwinBackbone(config=lowerCamelCase   )
       model.to(lowerCamelCase   )
       model.eval()
       UpperCAmelCase__							=    model(lowerCamelCase   )

       # verify feature maps
       self.parent.assertEqual(len(result.feature_maps   )			,   1   )
       self.parent.assertListEqual(list(result.feature_maps[0].shape   )			,   [self.batch_size, model.channels[-1], 4, 4]   )

       # verify channels
       self.parent.assertEqual(len(model.channels   )			,   1   )
      def 						UpperCAmelCase_    (      self :Optional[int]			,   lowerCamelCase :Any			,   lowerCamelCase :List[Any]			,   lowerCamelCase :Tuple   )  ->      Optional[Any]:
       UpperCAmelCase__							=    SwinForMaskedImageModeling(config=lowerCamelCase   )
       model.to(lowerCamelCase   )
       model.eval()
       UpperCAmelCase__							=    model(lowerCamelCase   )
       self.parent.assertEqual(
           result.logits.shape			,   (self.batch_size, self.num_channels, self.image_size, self.image_size)   )

       # test greyscale images
       UpperCAmelCase__							=    1
       UpperCAmelCase__							=    SwinForMaskedImageModeling(lowerCamelCase   )
       model.to(lowerCamelCase   )
       model.eval()

       UpperCAmelCase__							=    floats_tensor([self.batch_size, 1, self.image_size, self.image_size]   )
       UpperCAmelCase__							=    model(lowerCamelCase   )
       self.parent.assertEqual(result.logits.shape			,   (self.batch_size, 1, self.image_size, self.image_size)   )
      def 						UpperCAmelCase_    (      self :Union[str, Any]			,   lowerCamelCase :Any			,   lowerCamelCase :str			,   lowerCamelCase :List[Any]   )  ->      str:
       UpperCAmelCase__							=    self.type_sequence_label_size
       UpperCAmelCase__							=    SwinForImageClassification(lowerCamelCase   )
       model.to(lowerCamelCase   )
       model.eval()
       UpperCAmelCase__							=    model(lowerCamelCase			,   labels=lowerCamelCase   )
       self.parent.assertEqual(result.logits.shape			,   (self.batch_size, self.type_sequence_label_size)   )

       # test greyscale images
       UpperCAmelCase__							=    1
       UpperCAmelCase__							=    SwinForImageClassification(lowerCamelCase   )
       model.to(lowerCamelCase   )
       model.eval()

       UpperCAmelCase__							=    floats_tensor([self.batch_size, 1, self.image_size, self.image_size]   )
       UpperCAmelCase__							=    model(lowerCamelCase   )
       self.parent.assertEqual(result.logits.shape			,   (self.batch_size, self.type_sequence_label_size)   )







      def 						UpperCAmelCase_    (      self :List[str]   )  ->      Dict:
       UpperCAmelCase__							=    self.prepare_config_and_inputs()
       (
           (
           UpperCAmelCase__
       )			,			(
           UpperCAmelCase__
       )			,			(
           UpperCAmelCase__
       )			,			
       )							=    config_and_inputs
       UpperCAmelCase__							=    {"pixel_values": pixel_values}
       return config, inputs_dict



@require_torch
class       _UpperCamelCase   (					lowerCAmelCase						,							lowerCAmelCase						,							unittest.TestCase		):
      UpperCAmelCase_          =		(
          (
              SwinModel,
              SwinBackbone,
              SwinForImageClassification,
              SwinForMaskedImageModeling,
          )
          if is_torch_available()
          else ()
      )
      UpperCAmelCase_          =		(
          {"""feature-extraction""": SwinModel, """image-classification""": SwinForImageClassification}
          if is_torch_available()
          else {}
      )
      UpperCAmelCase_          =		True

      UpperCAmelCase_          =		False
      UpperCAmelCase_          =		False
      UpperCAmelCase_          =		False
      def 						UpperCAmelCase_    (      self :Any   )  ->      Dict:
       UpperCAmelCase__							=    SwinModelTester(self   )
       UpperCAmelCase__							=    ConfigTester(self			,   config_class=lowerCamelCase			,   embed_dim=37   )
      def 						UpperCAmelCase_    (      self :str   )  ->      str:
       self.create_and_test_config_common_properties()
       self.config_tester.create_and_test_config_to_json_string()
       self.config_tester.create_and_test_config_to_json_file()
       self.config_tester.create_and_test_config_from_and_save_pretrained()
       self.config_tester.create_and_test_config_with_num_labels()
       self.config_tester.check_config_can_be_init_without_params()
       self.config_tester.check_config_arguments_init()
      def 						UpperCAmelCase_    (      self :str   )  ->      Dict:
       return
      def 						UpperCAmelCase_    (      self :Tuple   )  ->      int:
       UpperCAmelCase__							=    self.model_tester.prepare_config_and_inputs()
       self.model_tester.create_and_check_model(*lowerCamelCase   )
      @unittest.skip(reason="Got `CUDA error: misaligned address` with PyTorch 2.0.0."   )
      def 						UpperCAmelCase_    (      self :Any   )  ->      List[str]:
       pass
      def 						UpperCAmelCase_    (      self :Union[str, Any]   )  ->      Optional[Any]:
       super().test_training_gradient_checkpointing()
      def 						UpperCAmelCase_    (      self :Optional[int]   )  ->      List[str]:
       UpperCAmelCase__							=    self.model_tester.prepare_config_and_inputs()
       self.model_tester.create_and_check_backbone(*lowerCamelCase   )
      def 						UpperCAmelCase_    (      self :List[str]   )  ->      Tuple:
       UpperCAmelCase__							=    self.model_tester.prepare_config_and_inputs()
       self.model_tester.create_and_check_for_masked_image_modeling(*lowerCamelCase   )
      def 						UpperCAmelCase_    (      self :Optional[Any]   )  ->      str:
       UpperCAmelCase__							=    self.model_tester.prepare_config_and_inputs()
       self.model_tester.create_and_check_for_image_classification(*lowerCamelCase   )
      @unittest.skip(reason="Swin does not use inputs_embeds"   )
      def 						UpperCAmelCase_    (      self :str   )  ->      Any:
       pass
      @unittest.skip(reason="Swin Transformer does not use feedforward chunking"   )
      def 						UpperCAmelCase_    (      self :Dict   )  ->      Union[str, Any]:
       pass
      def 						UpperCAmelCase_    (      self :Tuple   )  ->      Optional[Any]:
       UpperCAmelCase__			,			UpperCAmelCase__							=    self.model_tester.prepare_config_and_inputs_for_common()

       for model_class in self.all_model_classes:
        UpperCAmelCase__							=    model_class(lowerCamelCase   )
        self.assertIsInstance(model.get_input_embeddings()			,   (nn.Module)   )
        UpperCAmelCase__							=    model.get_output_embeddings()
        self.assertTrue(x is None or isinstance(lowerCamelCase			,   nn.Linear   )   )
      def 						UpperCAmelCase_    (      self :List[str]   )  ->      Union[str, Any]:
       UpperCAmelCase__			,			UpperCAmelCase__							=    self.model_tester.prepare_config_and_inputs_for_common()

       for model_class in self.all_model_classes:
        UpperCAmelCase__							=    model_class(lowerCamelCase   )
        UpperCAmelCase__							=    inspect.signature(model.forward   )
        # signature.parameters is an OrderedDict => so arg_names order is deterministic
        UpperCAmelCase__							=    [*signature.parameters.keys()]

        UpperCAmelCase__							=    ["pixel_values"]
        self.assertListEqual(arg_names[:1]			,   lowerCamelCase   )
      def 						UpperCAmelCase_    (      self :List[Any]   )  ->      str:
       UpperCAmelCase__			,			UpperCAmelCase__							=    self.model_tester.prepare_config_and_inputs_for_common()
       UpperCAmelCase__							=    True

       for model_class in self.all_model_classes:
        UpperCAmelCase__							=    True
        UpperCAmelCase__							=    False
        UpperCAmelCase__							=    True
        UpperCAmelCase__							=    model_class(lowerCamelCase   )
        model.to(lowerCamelCase   )
        model.eval()
        with torch.no_grad():
         UpperCAmelCase__							=    model(**self._prepare_for_class(lowerCamelCase			,   lowerCamelCase   )   )
        UpperCAmelCase__							=    outputs.attentions
        UpperCAmelCase__							=    len(self.model_tester.depths   )
        self.assertEqual(len(lowerCamelCase   )			,   lowerCamelCase   )

        # check that output_attentions also work using config
        del inputs_dict["output_attentions"]
        UpperCAmelCase__							=    True
        UpperCAmelCase__							=    config.window_size**2
        UpperCAmelCase__							=    model_class(lowerCamelCase   )
        model.to(lowerCamelCase   )
        model.eval()
        with torch.no_grad():
         UpperCAmelCase__							=    model(**self._prepare_for_class(lowerCamelCase			,   lowerCamelCase   )   )
        UpperCAmelCase__							=    outputs.attentions
        self.assertEqual(len(lowerCamelCase   )			,   lowerCamelCase   )

        self.assertListEqual(
            list(attentions[0].shape[-3:]   )			,   [self.model_tester.num_heads[0], window_size_squared, window_size_squared]			,   )
        UpperCAmelCase__							=    len(lowerCamelCase   )

        # Check attention is always last and order is fine
        UpperCAmelCase__							=    True
        UpperCAmelCase__							=    True
        UpperCAmelCase__							=    model_class(lowerCamelCase   )
        model.to(lowerCamelCase   )
        model.eval()
        with torch.no_grad():
         UpperCAmelCase__							=    model(**self._prepare_for_class(lowerCamelCase			,   lowerCamelCase   )   )

        # also another +1 for reshaped_hidden_states
        UpperCAmelCase__							=    1 if model_class.__name__ == "SwinBackbone" else 2
        self.assertEqual(out_len + added_hidden_states			,   len(lowerCamelCase   )   )

        UpperCAmelCase__							=    outputs.attentions

        self.assertEqual(len(lowerCamelCase   )			,   lowerCamelCase   )

        self.assertListEqual(
            list(self_attentions[0].shape[-3:]   )			,   [self.model_tester.num_heads[0], window_size_squared, window_size_squared]			,   )
      def 						UpperCAmelCase_    (      self :Optional[Any]			,   lowerCamelCase :List[Any]			,   lowerCamelCase :Tuple			,   lowerCamelCase :Optional[Any]			,   lowerCamelCase :Dict   )  ->      Optional[int]:
       UpperCAmelCase__							=    model_class(lowerCamelCase   )
       model.to(lowerCamelCase   )
       model.eval()

       with torch.no_grad():
        UpperCAmelCase__							=    model(**self._prepare_for_class(lowerCamelCase			,   lowerCamelCase   )   )

       UpperCAmelCase__							=    outputs.hidden_states

       UpperCAmelCase__							=    getattr(
           self.model_tester			,   "expected_num_hidden_layers"			,   len(self.model_tester.depths   ) + 1   )
       self.assertEqual(len(lowerCamelCase   )			,   lowerCamelCase   )

       # Swin has a different seq_length
       UpperCAmelCase__							=    (
           config.patch_size
           if isinstance(config.patch_size			,   collections.abc.Iterable   )
           else (config.patch_size, config.patch_size)
       )

       UpperCAmelCase__							=    (image_size[1] // patch_size[1]) * (image_size[0] // patch_size[0])

       self.assertListEqual(
           list(hidden_states[0].shape[-2:]   )			,   [num_patches, self.model_tester.embed_dim]			,   )

       if not model_class.__name__ == "SwinBackbone":
        UpperCAmelCase__							=    outputs.reshaped_hidden_states
        self.assertEqual(len(lowerCamelCase   )			,   lowerCamelCase   )

        UpperCAmelCase__			,			UpperCAmelCase__			,			UpperCAmelCase__			,			UpperCAmelCase__							=    reshaped_hidden_states[0].shape
        UpperCAmelCase__							=    (
            reshaped_hidden_states[0].view(lowerCamelCase			,   lowerCamelCase			,   height * width   ).permute(0			,   2			,   1   )
        )
        self.assertListEqual(
            list(reshaped_hidden_states.shape[-2:]   )			,   [num_patches, self.model_tester.embed_dim]			,   )
      def 						UpperCAmelCase_    (      self :Any   )  ->      Dict:
       UpperCAmelCase__			,			UpperCAmelCase__							=    self.model_tester.prepare_config_and_inputs_for_common()

       UpperCAmelCase__							=    (
           self.model_tester.image_size
           if isinstance(self.model_tester.image_size			,   collections.abc.Iterable   )
           else (self.model_tester.image_size, self.model_tester.image_size)
       )

       for model_class in self.all_model_classes:
        UpperCAmelCase__							=    True
        self.check_hidden_states_output(lowerCamelCase			,   lowerCamelCase			,   lowerCamelCase			,   lowerCamelCase   )

        # check that output_hidden_states also work using config
        del inputs_dict["output_hidden_states"]
        UpperCAmelCase__							=    True

        self.check_hidden_states_output(lowerCamelCase			,   lowerCamelCase			,   lowerCamelCase			,   lowerCamelCase   )
      def 						UpperCAmelCase_    (      self :List[str]   )  ->      List[str]:
       UpperCAmelCase__			,			UpperCAmelCase__							=    self.model_tester.prepare_config_and_inputs_for_common()
       UpperCAmelCase__							=    3

       UpperCAmelCase__							=    (
           self.model_tester.image_size
           if isinstance(self.model_tester.image_size			,   collections.abc.Iterable   )
           else (self.model_tester.image_size, self.model_tester.image_size)
       )
       UpperCAmelCase__							=    (
           config.patch_size
           if isinstance(config.patch_size			,   collections.abc.Iterable   )
           else (config.patch_size, config.patch_size)
       )

       UpperCAmelCase__							=    image_size[0] + patch_size[0] - (image_size[0] % patch_size[0])
       UpperCAmelCase__							=    image_size[1] + patch_size[1] - (image_size[1] % patch_size[1])

       for model_class in self.all_model_classes:
        UpperCAmelCase__							=    True
        self.check_hidden_states_output(lowerCamelCase			,   lowerCamelCase			,   lowerCamelCase			,   (padded_height, padded_width)   )

        # check that output_hidden_states also work using config
        del inputs_dict["output_hidden_states"]
        UpperCAmelCase__							=    True
        self.check_hidden_states_output(lowerCamelCase			,   lowerCamelCase			,   lowerCamelCase			,   (padded_height, padded_width)   )
      @slow
      def 						UpperCAmelCase_    (      self :Any   )  ->      List[Any]:
       for model_name in SWIN_PRETRAINED_MODEL_ARCHIVE_LIST[:1]:
        UpperCAmelCase__							=    SwinModel.from_pretrained(lowerCamelCase   )
        self.assertIsNotNone(lowerCamelCase   )







      def 						UpperCAmelCase_    (      self :List[Any]   )  ->      Optional[Any]:
       UpperCAmelCase__			,			UpperCAmelCase__							=    self.model_tester.prepare_config_and_inputs_for_common()

       UpperCAmelCase__							=    _config_zero_init(lowerCamelCase   )
       for model_class in self.all_model_classes:
        UpperCAmelCase__							=    model_class(config=lowerCamelCase   )
        for name, param in model.named_parameters():
         if "embeddings" not in name and param.requires_grad:
          self.assertIn(
              ((param.data.mean() * 1e9).round() / 1e9).item()			,   [0.0, 1.0]			,   msg=f'''Parameter {name} of model {model_class} seems not properly initialized'''			,   )



@require_vision
@require_torch
class       _UpperCamelCase   (					unittest.TestCase		):
      @cached_property
      def 						UpperCAmelCase_    (      self :Optional[int]   )  ->      int:
       return (
           AutoImageProcessor.from_pretrained("microsoft/swin-tiny-patch4-window7-224"   )
           if is_vision_available()
           else None
       )







      @slow
      def 						UpperCAmelCase_    (      self :List[Any]   )  ->      Any:
       UpperCAmelCase__							=    SwinForImageClassification.from_pretrained("microsoft/swin-tiny-patch4-window7-224"   ).to(lowerCamelCase   )
       UpperCAmelCase__							=    self.default_image_processor

       UpperCAmelCase__							=    Image.open("./tests/fixtures/tests_samples/COCO/000000039769.png"   )
       UpperCAmelCase__							=    image_processor(images=lowerCamelCase			,   return_tensors="pt"   ).to(lowerCamelCase   )

       # forward pass
       with torch.no_grad():
        UpperCAmelCase__							=    model(**lowerCamelCase   )

       # verify the logits
       UpperCAmelCase__							=    torch.Size((1, 1000)   )
       self.assertEqual(outputs.logits.shape			,   lowerCamelCase   )
       UpperCAmelCase__							=    torch.tensor([-0.09_48, -0.64_54, -0.09_21]   ).to(lowerCamelCase   )
       self.assertTrue(torch.allclose(outputs.logits[0, :3]			,   lowerCamelCase			,   atol=1e-4   )   )



@require_torch
class       _UpperCamelCase   (					unittest.TestCase						,							lowerCAmelCase		):
      UpperCAmelCase_          =		(SwinBackbone,) if is_torch_available() else ()
      UpperCAmelCase_          =		SwinConfig
      def 						UpperCAmelCase_    (      self :Tuple   )  ->      List[Any]:
       UpperCAmelCase__							=    SwinModelTester(self   )











import random
from typing import List, Optional, Union

import numpy as np
import torch

from ..state import AcceleratorState
from .constants import CUDA_DISTRIBUTED_TYPES
from .dataclasses import DistributedType, RNGType
from .imports import is_npu_available, is_tpu_available, is_xpu_available


if is_tpu_available(check_device=False):
  import torch_xla.core.xla_model as xm



def     lowerCAmelCase		(		_lowerCAmelCase				: int    ,				_lowerCAmelCase				: bool = False						):




 """simple docstring"""







 if device_specific:
  seed += AcceleratorState().process_index
 random.seed(_lowerCAmelCase						)
 np.random.seed(_lowerCAmelCase						)
 torch.manual_seed(_lowerCAmelCase						)
 if is_xpu_available():
  torch.xpu.manual_seed_all(_lowerCAmelCase						)
 elif is_npu_available():
  torch.npu.manual_seed_all(_lowerCAmelCase						)
 else:
  torch.cuda.manual_seed_all(_lowerCAmelCase						)
 # ^^ safe to call this function even if cuda is not available
 if is_tpu_available():
  xm.set_rng_state(_lowerCAmelCase						)



def     lowerCAmelCase		(		_lowerCAmelCase				: Optional[RNGType] = None    ,				_lowerCAmelCase				: Optional[torch.Generator] = None						):




 """simple docstring"""







 if rng_type == RNGType.TORCH:
  UpperCAmelCase__							=    torch.get_rng_state()
 elif rng_type == RNGType.CUDA:
  UpperCAmelCase__							=    torch.cuda.get_rng_state()
 elif rng_type == RNGType.XLA:
  assert is_tpu_available(), "Can't synchronize XLA seeds on an environment without TPUs."
  UpperCAmelCase__							=    torch.tensor(xm.get_rng_state()						)
 elif rng_type == RNGType.NPU:
  assert is_npu_available(), "Can't synchronize NPU seeds on an environment without NPUs."
  UpperCAmelCase__							=    torch.npu.get_rng_state()
 elif rng_type == RNGType.XPU:
  assert is_xpu_available(), "Can't synchronize XPU seeds on an environment without XPUs."
  UpperCAmelCase__							=    torch.xpu.get_rng_state()
 elif rng_type == RNGType.GENERATOR:
  assert generator is not None, "Need a generator to synchronize its seed."
  UpperCAmelCase__							=    generator.get_state()

 # Broadcast the rng state from device 0 to other devices
 UpperCAmelCase__							=    AcceleratorState()
 if state.distributed_type == DistributedType.TPU:
  UpperCAmelCase__							=    rng_state.to(xm.xla_device()						)
  xm.collective_broadcast([rng_state]						)
  xm.mark_step()
  UpperCAmelCase__							=    rng_state.cpu()
 elif (
     state.distributed_type in CUDA_DISTRIBUTED_TYPES
     or state.distributed_type == DistributedType.MULTI_NPU
     or state.distributed_type == DistributedType.MULTI_XPU
 ):
  UpperCAmelCase__							=    rng_state.to(state.device						)
  torch.distributed.broadcast(_lowerCAmelCase    ,				0						)
  UpperCAmelCase__							=    rng_state.cpu()
 elif state.distributed_type == DistributedType.MULTI_CPU:
  torch.distributed.broadcast(_lowerCAmelCase    ,				0						)

 # Set the broadcast rng state
 if rng_type == RNGType.TORCH:
  torch.set_rng_state(_lowerCAmelCase						)
 elif rng_type == RNGType.CUDA:
  torch.cuda.set_rng_state(_lowerCAmelCase						)
 elif rng_type == RNGType.NPU:
  torch.npu.set_rng_state(_lowerCAmelCase						)
 elif rng_type == RNGType.XPU:
  torch.xpu.set_rng_state(_lowerCAmelCase						)
 elif rng_type == RNGType.XLA:
  xm.set_rng_state(rng_state.item()						)
 elif rng_type == RNGType.GENERATOR:
  generator.set_state(_lowerCAmelCase						)



def     lowerCAmelCase		(		_lowerCAmelCase				: List[Union[str, RNGType]]    ,				_lowerCAmelCase				: Optional[torch.Generator] = None						):




 """simple docstring"""







 for rng_type in rng_types:
  synchronize_rng_state(RNGType(_lowerCAmelCase						)    ,				generator=_lowerCAmelCase						)





