



'''simple docstring'''


import json
import os
import tempfile
from unittest.mock import patch

import torch
from torch.utils.data import DataLoader, TensorDataset

from accelerate import DistributedType, infer_auto_device_map, init_empty_weights
from accelerate.accelerator import Accelerator
from accelerate.state import GradientState, PartialState
from accelerate.test_utils import require_bnb, require_multi_gpu, slow
from accelerate.test_utils.testing import AccelerateTestCase, require_cuda
from accelerate.utils import patch_environment







def 				__lowerCamelCase				(   )       ->	List[Any]:






   """simple docstring"""
   UpperCamelCase					=  torch.nn.Linear(2    ,						4    )
   UpperCamelCase					=  torch.optim.AdamW(model.parameters()    ,						lr=1.0    )
   UpperCamelCase					=  torch.optim.lr_scheduler.OneCycleLR(A__    ,						max_lr=0.01    ,						steps_per_epoch=2    ,						epochs=1    )
   UpperCamelCase					=  DataLoader(TensorDataset(torch.tensor([1, 2, 3]    )    )    )
   UpperCamelCase					=  DataLoader(TensorDataset(torch.tensor([4, 5, 6]    )    )    )

   return model, optimizer, scheduler, train_dl, valid_dl







def 				__lowerCamelCase				(   A__    )       ->	str:






   """simple docstring"""
   return (model.weight.abs().sum() + model.bias.abs().sum()).item()







def 				__lowerCamelCase				(   A__    )       ->	Dict:






   """simple docstring"""
   UpperCamelCase					=  torch.nn.Linear(*tuple(model.weight.T.shape    )    ).state_dict()
   model.load_state_dict(A__    )

class  SCREAMING_SNAKE_CASE (  _a		):






    """simple docstring"""





    @require_cuda
    def  A       (  self      :					Any ):


       """simple docstring"""







       UpperCamelCase					=  Accelerator()
       assert PartialState._shared_state["_cpu"] is False
       assert PartialState._shared_state["device"].type == "cuda"
       with self.assertRaises(UpperCamelCase__ ):
          UpperCamelCase					=  Accelerator(cpu=UpperCamelCase__ )






    def  A       (  self      :					Any ):


       """simple docstring"""







       UpperCamelCase					=  Accelerator()
       UpperCamelCase					=  GradientState()
       assert state.num_steps == 1
       UpperCamelCase					=  4
       assert state.num_steps == 4

       assert state.sync_gradients is True
       UpperCamelCase					=  False
       assert state.sync_gradients is False
       GradientState._reset_state()






    def  A       (  self      :					List[str] ):


       """simple docstring"""







       UpperCamelCase					=  Accelerator()
       UpperCamelCase				,      UpperCamelCase				,      UpperCamelCase				,      UpperCamelCase				,      UpperCamelCase					=  create_components()

       (
           (
           UpperCamelCase
       )				,      (
           UpperCamelCase
       )				,      (
           UpperCamelCase
       )				,      (
           UpperCamelCase
       )				,      (
           UpperCamelCase
       )				,      
       )					=  accelerator.prepare(UpperCamelCase__  ,						UpperCamelCase__  ,						UpperCamelCase__  ,						UpperCamelCase__  ,						UpperCamelCase__ )

       self.assertTrue(prepared_model in accelerator._models )
       self.assertTrue(prepared_optimizer in accelerator._optimizers )
       self.assertTrue(prepared_scheduler in accelerator._schedulers )
       self.assertTrue(prepared_train_dl in accelerator._dataloaders )
       self.assertTrue(prepared_valid_dl in accelerator._dataloaders )






    def  A       (  self      :					List[str] ):


       """simple docstring"""







       UpperCamelCase					=  Accelerator()
       UpperCamelCase				,      UpperCamelCase				,      UpperCamelCase				,      UpperCamelCase				,      UpperCamelCase					=  create_components()
       accelerator.prepare(UpperCamelCase__  ,						UpperCamelCase__  ,						UpperCamelCase__  ,						UpperCamelCase__  ,						UpperCamelCase__ )
       accelerator.free_memory()

       self.assertTrue(len(accelerator._models ) == 0 )
       self.assertTrue(len(accelerator._optimizers ) == 0 )
       self.assertTrue(len(accelerator._schedulers ) == 0 )
       self.assertTrue(len(accelerator._dataloaders ) == 0 )






    def  A       (  self      :					Optional[Any] ):


       """simple docstring"""







       PartialState._reset_state()

       # Mock torch.cuda.set_device to avoid an exception as the device doesn't exist
       def noop(*UpperCamelCase__      :					Dict  ,						**UpperCamelCase__      :					int ):
          pass

       with patch('torch.cuda.set_device'  ,						UpperCamelCase__ ), patch_environment(ACCELERATE_TORCH_DEVICE='cuda:64' ):
          UpperCamelCase					=  Accelerator()
          self.assertEqual(str(accelerator.state.device )  ,						'cuda:64' )






    def  A       (  self      :					Optional[int] ):


       """simple docstring"""







       UpperCamelCase					=  Accelerator()
       UpperCamelCase				,      UpperCamelCase				,      UpperCamelCase				,      UpperCamelCase				,      UpperCamelCase					=  create_components()
       accelerator.prepare(UpperCamelCase__  ,						UpperCamelCase__  ,						UpperCamelCase__  ,						UpperCamelCase__  ,						UpperCamelCase__ )

       UpperCamelCase					=  get_signature(UpperCamelCase__ )

       with tempfile.TemporaryDirectory() as tmpdirname:
          accelerator.save_state(UpperCamelCase__ )

          # make sure random weights don't match
          load_random_weights(UpperCamelCase__ )
          self.assertTrue(abs(model_signature - get_signature(UpperCamelCase__ ) ) > 1E-3 )

          # make sure loaded weights match
          accelerator.load_state(UpperCamelCase__ )
          self.assertTrue(abs(model_signature - get_signature(UpperCamelCase__ ) ) < 1E-3 )






    def  A       (  self      :					List[Any] ):


       """simple docstring"""







       UpperCamelCase					=  Accelerator()
       UpperCamelCase				,      UpperCamelCase				,      UpperCamelCase				,      UpperCamelCase				,      UpperCamelCase					=  create_components()
       accelerator.prepare(UpperCamelCase__  ,						UpperCamelCase__  ,						UpperCamelCase__  ,						UpperCamelCase__  ,						UpperCamelCase__ )

       UpperCamelCase					=  get_signature(UpperCamelCase__ )

       # saving hook
       def save_config(UpperCamelCase__      :					str  ,						UpperCamelCase__      :					Union[str, Any]  ,						UpperCamelCase__      :					str ):
          UpperCamelCase					=  {'class_name': models[0].__class__.__name__}

          with open(os.path.join(UpperCamelCase__  ,						'data.json' )  ,						'w' ) as f:
             json.dump(UpperCamelCase__  ,						UpperCamelCase__ )

        # loading hook
       def load_config(UpperCamelCase__      :					Tuple  ,						UpperCamelCase__      :					Union[str, Any] ):
          with open(os.path.join(UpperCamelCase__  ,						'data.json' )  ,						'r' ) as f:
             UpperCamelCase					=  json.load(UpperCamelCase__ )

          UpperCamelCase					=  config['class_name']

       UpperCamelCase					=  accelerator.register_save_state_pre_hook(UpperCamelCase__ )
       UpperCamelCase					=  accelerator.register_load_state_pre_hook(UpperCamelCase__ )

       with tempfile.TemporaryDirectory() as tmpdirname:
          accelerator.save_state(UpperCamelCase__ )

          # make sure random weights don't match with hooks
          load_random_weights(UpperCamelCase__ )
          self.assertTrue(abs(model_signature - get_signature(UpperCamelCase__ ) ) > 1E-3 )

          # random class name to verify correct one is loaded
          UpperCamelCase					=  'random'

          # make sure loaded weights match with hooks
          accelerator.load_state(UpperCamelCase__ )
          self.assertTrue(abs(model_signature - get_signature(UpperCamelCase__ ) ) < 1E-3 )

          # mode.class_name is loaded from config
          self.assertTrue(model.class_name == model.__class__.__name__ )

       # remove hooks
       save_hook.remove()
       load_hook.remove()

       with tempfile.TemporaryDirectory() as tmpdirname:
          accelerator.save_state(UpperCamelCase__ )

          # make sure random weights don't match with hooks removed
          load_random_weights(UpperCamelCase__ )
          self.assertTrue(abs(model_signature - get_signature(UpperCamelCase__ ) ) > 1E-3 )

          # random class name to verify correct one is loaded
          UpperCamelCase					=  'random'

          # make sure loaded weights match with hooks removed
          accelerator.load_state(UpperCamelCase__ )
          self.assertTrue(abs(model_signature - get_signature(UpperCamelCase__ ) ) < 1E-3 )

          # mode.class_name is NOT loaded from config
          self.assertTrue(model.class_name != model.__class__.__name__ )






    def  A       (  self      :					str ):


       """simple docstring"""







       UpperCamelCase					=  Accelerator()
       UpperCamelCase				,      UpperCamelCase				,      UpperCamelCase				,      UpperCamelCase				,      UpperCamelCase					=  create_components()
       UpperCamelCase					=  None

       # This should work
       UpperCamelCase				,      UpperCamelCase				,      UpperCamelCase				,      UpperCamelCase				,      UpperCamelCase				,      UpperCamelCase					=  accelerator.prepare(
           UpperCamelCase__  ,						UpperCamelCase__  ,						UpperCamelCase__  ,						UpperCamelCase__  ,						UpperCamelCase__  ,						UpperCamelCase__ )
       self.assertTrue(dummy_obj is None )






    def  A       (  self      :					Any ):


       """simple docstring"""







       UpperCamelCase					=  Accelerator()
       UpperCamelCase				,      UpperCamelCase				,      UpperCamelCase				,      UpperCamelCase				,      UpperCamelCase					=  create_components()
       UpperCamelCase					=  [1, 2, 3]

       # This should work
       UpperCamelCase				,      UpperCamelCase				,      UpperCamelCase				,      UpperCamelCase				,      UpperCamelCase				,      UpperCamelCase					=  accelerator.prepare(
           UpperCamelCase__  ,						UpperCamelCase__  ,						UpperCamelCase__  ,						UpperCamelCase__  ,						UpperCamelCase__  ,						UpperCamelCase__ )
       self.assertEqual(
           getattr(UpperCamelCase__  ,						'_is_accelerate_prepared'  ,						UpperCamelCase__ )  ,						UpperCamelCase__  ,						'Dummy object should have `_is_accelerate_prepared` set to `True`'  ,						)
       self.assertEqual(
           getattr(UpperCamelCase__  ,						'_is_accelerate_prepared'  ,						UpperCamelCase__ )  ,						UpperCamelCase__  ,						'Model is missing `_is_accelerator_prepared` or is set to `False`'  ,						)
       self.assertEqual(
           getattr(UpperCamelCase__  ,						'_is_accelerate_prepared'  ,						UpperCamelCase__ )  ,						UpperCamelCase__  ,						'Optimizer is missing `_is_accelerator_prepared` or is set to `False`'  ,						)
       self.assertEqual(
           getattr(UpperCamelCase__  ,						'_is_accelerate_prepared'  ,						UpperCamelCase__ )  ,						UpperCamelCase__  ,						'Scheduler is missing `_is_accelerator_prepared` or is set to `False`'  ,						)
       self.assertEqual(
           getattr(UpperCamelCase__  ,						'_is_accelerate_prepared'  ,						UpperCamelCase__ )  ,						UpperCamelCase__  ,						'Train Dataloader is missing `_is_accelerator_prepared` or is set to `False`'  ,						)
       self.assertEqual(
           getattr(UpperCamelCase__  ,						'_is_accelerate_prepared'  ,						UpperCamelCase__ )  ,						UpperCamelCase__  ,						'Valid Dataloader is missing `_is_accelerator_prepared` or is set to `False`'  ,						)






    @slow
    @require_bnb
    def  A       (  self      :					List[str] ):


       """simple docstring"""







       from transformers import AutoModelForCausalLM

       UpperCamelCase					=  AutoModelForCausalLM.from_pretrained(
           'EleutherAI/gpt-neo-125m'  ,						load_in_abit=UpperCamelCase__  ,						device_map={'': 0}  ,						)
       UpperCamelCase					=  Accelerator()

       # This should work
       UpperCamelCase					=  accelerator.prepare(UpperCamelCase__ )






    @slow
    @require_bnb
    def  A       (  self      :					Tuple ):


       """simple docstring"""







       from transformers import AutoModelForCausalLM

       UpperCamelCase					=  Accelerator()

       with init_empty_weights():
          UpperCamelCase					=  AutoModelForCausalLM.from_pretrained(
              'EleutherAI/gpt-neo-125m'  ,						)
          model.tie_weights()
          UpperCamelCase					=  infer_auto_device_map(UpperCamelCase__ )
          UpperCamelCase					=  'cpu'

       UpperCamelCase					=  AutoModelForCausalLM.from_pretrained(
           'EleutherAI/gpt-neo-125m'  ,						device_map=UpperCamelCase__  ,						load_in_abit=UpperCamelCase__  ,						llm_inta_enable_fpaa_cpu_offload=UpperCamelCase__ )

       # This should not work and get value error
       with self.assertRaises(UpperCamelCase__ ):
          UpperCamelCase					=  accelerator.prepare(UpperCamelCase__ )






    @slow
    @require_bnb
    @require_multi_gpu
    def  A       (  self      :					Optional[Any] ):


       """simple docstring"""







       from transformers import AutoModelForCausalLM

       UpperCamelCase					=  {'distributed_type': DistributedType.MULTI_GPU}

       with init_empty_weights():
          UpperCamelCase					=  AutoModelForCausalLM.from_pretrained(
              'EleutherAI/gpt-neo-125m'  ,						)
          model.tie_weights()
          UpperCamelCase					=  infer_auto_device_map(UpperCamelCase__ )
          UpperCamelCase					=  1

       UpperCamelCase					=  AutoModelForCausalLM.from_pretrained(
           'EleutherAI/gpt-neo-125m'  ,						load_in_abit=UpperCamelCase__  ,						device_map=UpperCamelCase__  ,						)
       UpperCamelCase					=  Accelerator()

       # This should not work and get value error
       with self.assertRaises(UpperCamelCase__ ):
          UpperCamelCase					=  accelerator.prepare(UpperCamelCase__ )

       PartialState._reset_state()






    @slow
    @require_bnb
    @require_multi_gpu
    def  A       (  self      :					Optional[Any] ):


       """simple docstring"""







       from transformers import AutoModelForCausalLM

       with init_empty_weights():
          UpperCamelCase					=  AutoModelForCausalLM.from_pretrained(
              'EleutherAI/gpt-neo-125m'  ,						)
          UpperCamelCase					=  infer_auto_device_map(UpperCamelCase__ )
          UpperCamelCase					=  1

       UpperCamelCase					=  AutoModelForCausalLM.from_pretrained(
           'EleutherAI/gpt-neo-125m'  ,						load_in_abit=UpperCamelCase__  ,						device_map=UpperCamelCase__  ,						)
       UpperCamelCase					=  Accelerator()

       # This should work
       UpperCamelCase					=  accelerator.prepare(UpperCamelCase__ )






    @require_cuda
    def  A       (  self      :					Union[str, Any] ):


       """simple docstring"""







       UpperCamelCase					=  torch.nn.Linear(1_0  ,						1_0 )
       UpperCamelCase					=  torch.optim.SGD(model.parameters()  ,						lr=0.0_1 )
       UpperCamelCase					=  Accelerator(cpu=UpperCamelCase__ )
       UpperCamelCase					=  accelerator.prepare(UpperCamelCase__ )






'''simple docstring'''


from typing import Optional, Tuple

import jax
import jax.numpy as jnp
from flax import linen as nn
from flax.core.frozen_dict import FrozenDict
from transformers import CLIPConfig, FlaxPreTrainedModel
from transformers.models.clip.modeling_flax_clip import FlaxCLIPVisionModule







def 				__lowerCamelCase				(   A__    ,						A__    ,						A__=1e-1_2    )       ->	Dict:






   """simple docstring"""
   UpperCamelCase					=  jnp.divide(emb_a.T    ,						jnp.clip(jnp.linalg.norm(A__    ,						axis=1    )    ,						a_min=A__    )    ).T
   UpperCamelCase					=  jnp.divide(emb_a.T    ,						jnp.clip(jnp.linalg.norm(A__    ,						axis=1    )    ,						a_min=A__    )    ).T
   return jnp.matmul(A__    ,						norm_emb_a.T    )

class  SCREAMING_SNAKE_CASE (  nn.Module		):






    """simple docstring"""





    _SCREAMING_SNAKE_CASE	  =  42
    _SCREAMING_SNAKE_CASE	  =  jnp.floataa






    def  A       (  self      :					List[Any] ):


       """simple docstring"""







       UpperCamelCase					=  FlaxCLIPVisionModule(self.config.vision_config )
       UpperCamelCase					=  nn.Dense(self.config.projection_dim  ,						use_bias=UpperCamelCase__  ,						dtype=self.dtype )

       UpperCamelCase					=  self.param('concept_embeds'  ,						jax.nn.initializers.ones  ,						(1_7, self.config.projection_dim) )
       UpperCamelCase					=  self.param(
           'special_care_embeds'  ,						jax.nn.initializers.ones  ,						(3, self.config.projection_dim) )

       UpperCamelCase					=  self.param('concept_embeds_weights'  ,						jax.nn.initializers.ones  ,						(1_7,) )
       UpperCamelCase					=  self.param('special_care_embeds_weights'  ,						jax.nn.initializers.ones  ,						(3,) )






    def __call__(  self      :					str  ,						UpperCamelCase__      :					List[str] ):


       """simple docstring"""







       UpperCamelCase					=  self.vision_model(UpperCamelCase__ )[1]
       UpperCamelCase					=  self.visual_projection(UpperCamelCase__ )

       UpperCamelCase					=  jax_cosine_distance(UpperCamelCase__  ,						self.special_care_embeds )
       UpperCamelCase					=  jax_cosine_distance(UpperCamelCase__  ,						self.concept_embeds )

       # increase this value to create a stronger `nfsw` filter
       # at the cost of increasing the possibility of filtering benign image inputs
       UpperCamelCase					=  0.0

       UpperCamelCase					=  special_cos_dist - self.special_care_embeds_weights[None, :] + adjustment
       UpperCamelCase					=  jnp.round(UpperCamelCase__  ,						3 )
       UpperCamelCase					=  jnp.any(special_scores > 0  ,						axis=1  ,						keepdims=UpperCamelCase__ )
       # Use a lower threshold if an image has any special care concept
       UpperCamelCase					=  is_special_care * 0.0_1

       UpperCamelCase					=  cos_dist - self.concept_embeds_weights[None, :] + special_adjustment
       UpperCamelCase					=  jnp.round(UpperCamelCase__  ,						3 )
       UpperCamelCase					=  jnp.any(concept_scores > 0  ,						axis=1 )

       return has_nsfw_concepts

class  SCREAMING_SNAKE_CASE (  _a		):






    """simple docstring"""





    _SCREAMING_SNAKE_CASE	  =  CLIPConfig
    _SCREAMING_SNAKE_CASE	  =  """clip_input"""
    _SCREAMING_SNAKE_CASE	  =  FlaxStableDiffusionSafetyCheckerModule






    def __init__(  self      :					Union[str, Any]  ,						UpperCamelCase__      :					CLIPConfig  ,						UpperCamelCase__      :					Optional[Tuple] = None  ,						UpperCamelCase__      :					int = 0  ,						UpperCamelCase__      :					jnp.dtype = jnp.floataa  ,						UpperCamelCase__      :					bool = True  ,						**UpperCamelCase__      :					List[str]  ,						):


       """simple docstring"""







       if input_shape is None:
          UpperCamelCase					=  (1, 2_2_4, 2_2_4, 3)
       UpperCamelCase					=  self.module_class(config=UpperCamelCase__  ,						dtype=UpperCamelCase__  ,						**UpperCamelCase__ )
       super().__init__(UpperCamelCase__  ,						UpperCamelCase__  ,						input_shape=UpperCamelCase__  ,						seed=UpperCamelCase__  ,						dtype=UpperCamelCase__  ,						_do_init=_do_init )






    def  A       (  self      :					int  ,						UpperCamelCase__      :					jax.random.KeyArray  ,						UpperCamelCase__      :					Tuple  ,						UpperCamelCase__      :					FrozenDict = None ):


       """simple docstring"""







       UpperCamelCase					=  jax.random.normal(UpperCamelCase__  ,						UpperCamelCase__ )

       UpperCamelCase				,      UpperCamelCase					=  jax.random.split(UpperCamelCase__ )
       UpperCamelCase					=  {'params': params_rng, 'dropout': dropout_rng}

       UpperCamelCase					=  self.module.init(UpperCamelCase__  ,						UpperCamelCase__ )['params']

       return random_params






    def __call__(  self      :					List[Any]  ,						UpperCamelCase__      :					Dict  ,						UpperCamelCase__      :					dict = None  ,						):


       """simple docstring"""







       UpperCamelCase					=  jnp.transpose(UpperCamelCase__  ,						(0, 2, 3, 1) )

       return self.module.apply(
           {'params': params or self.params}  ,						jnp.array(UpperCamelCase__  ,						dtype=jnp.floataa )  ,						rngs={}  ,						)

