code: 
  
  
  import inspect
  import warnings
  
  import torch
  
  from .state import AcceleratorState, GradientState
  from .utils import DistributedType, honor_type, is_tpu_available
  
  
  if is_tpu_available(check_device=False):
  		import torch_xla.core.xla_model as xm
  
  
  
  
  
  
  
  def 							UpperCAmelCase				(     UpperCAmelCase				,UpperCAmelCase      )->		Dict:
  
  
  
  
  
  
  
  					'''simple docstring'''
  
  
  
  					if isinstance(UpperCAmelCase				,(list, tuple)      ):
  										return honor_type(UpperCAmelCase				,(move_to_device(UpperCAmelCase				,UpperCAmelCase      ) for t in state)      )
  					elif isinstance(UpperCAmelCase				,UpperCAmelCase      ):
  										return type(UpperCAmelCase      )({k: move_to_device(UpperCAmelCase				,UpperCAmelCase      ) for k, v in state.items()}      )
  					elif isinstance(UpperCAmelCase				,torch.Tensor      ):
  										return state.to(UpperCAmelCase      )
  					return state
  
  
  
  
  
  
  class 						snake_case			(       torch.optim.Optimizer	):
  			'''simple docstring'''
  
  			def __init__(	self     :     Optional[int]							,							lowerCAmelCase_     :     Union[str, Any]							,							lowerCAmelCase_     :     int=True							,							lowerCAmelCase_     :     List[Any]=None					)					->					List[Any]:
  
  
  
  
  								"""simple docstring"""
  
  
  
  
  
  
  								SCREAMING_SNAKE_CASE_	      =       optimizer
  								SCREAMING_SNAKE_CASE_	      =       scaler
  								SCREAMING_SNAKE_CASE_	      =       AcceleratorState()
  								SCREAMING_SNAKE_CASE_	      =       GradientState()
  								SCREAMING_SNAKE_CASE_	      =       device_placement
  								SCREAMING_SNAKE_CASE_	      =       False
  								SCREAMING_SNAKE_CASE_	      =       None
  
  								# Handle device placement
  								if device_placement:
  													SCREAMING_SNAKE_CASE_	      =       self.optimizer.state_dict()
  													if self.accelerator_state.distributed_type == DistributedType.TPU:
  																		xm.send_cpu_data_to_device(__A							,							self.accelerator_state.device					)
  													else:
  																		SCREAMING_SNAKE_CASE_	      =       move_to_device(__A							,							self.accelerator_state.device					)
  													self.optimizer.load_state_dict(__A					)
  
  			@property
  			def _lowercase (	self     :     Any					)					->					Dict:
  
  
  
  
  								"""simple docstring"""
  
  
  
  
  
  
  								return self.optimizer.state
  
  			@state.setter
  			def _lowercase (	self     :     Dict							,							lowerCAmelCase_     :     int					)					->					Optional[Any]:
  
  
  
  
  								"""simple docstring"""
  
  
  
  
  
  
  								SCREAMING_SNAKE_CASE_	      =       state
  
  			@property
  			def _lowercase (	self     :     Optional[Any]					)					->					List[str]:
  
  
  
  
  								"""simple docstring"""
  
  
  
  
  
  
  								return self.optimizer.param_groups
  
  			@param_groups.setter
  			def _lowercase (	self     :     Optional[Any]							,							lowerCAmelCase_     :     Optional[Any]					)					->					List[Any]:
  
  
  
  
  								"""simple docstring"""
  
  
  
  
  
  
  								SCREAMING_SNAKE_CASE_	      =       param_groups
  
  			@property
  			def _lowercase (	self     :     Any					)					->					List[str]:
  
  
  
  
  								"""simple docstring"""
  
  
  
  
  
  
  								return self.optimizer.defaults
  
  			@defaults.setter
  			def _lowercase (	self     :     int							,							lowerCAmelCase_     :     List[Any]					)					->					int:
  
  
  
  
  								"""simple docstring"""
  
  
  
  
  
  
  								SCREAMING_SNAKE_CASE_	      =       defaults
  
  			def _lowercase (	self     :     Optional[int]							,							lowerCAmelCase_     :     Tuple					)					->					Dict:
  
  
  
  
  								"""simple docstring"""
  
  
  
  
  
  
  								self.optimizer.add_param_group(__A					)
  
  			def _lowercase (	self     :     Any							,							lowerCAmelCase_     :     Optional[int]					)					->					Union[str, Any]:
  
  
  
  
  								"""simple docstring"""
  
  
  
  
  
  
  								if self.accelerator_state.distributed_type == DistributedType.TPU and self.device_placement:
  													xm.send_cpu_data_to_device(__A							,							self.accelerator_state.device					)
  								self.optimizer.load_state_dict(__A					)
  
  			def _lowercase (	self     :     Any					)					->					Tuple:
  
  
  
  
  								"""simple docstring"""
  
  
  
  
  
  
  								return self.optimizer.state_dict()
  
  			def _lowercase (	self     :     Optional[int]							,							lowerCAmelCase_     :     int=None					)					->					Any:
  
  
  
  
  								"""simple docstring"""
  
  
  
  
  
  
  								if self.gradient_state.sync_gradients:
  													SCREAMING_SNAKE_CASE_	      =       "set_to_none" in inspect.signature(self.optimizer.zero_grad					).parameters
  													if accept_arg:
  																		if set_to_none is None:
  																							SCREAMING_SNAKE_CASE_	      =       False
  																		self.optimizer.zero_grad(set_to_none=__A					)
  													else:
  																		if set_to_none is not None:
  																							raise ValueError('''`set_to_none` for Optimizer.zero_grad` is not supported by this optimizer.'''					)
  																		self.optimizer.zero_grad()
  
  			def _lowercase (	self     :     str							,							lowerCAmelCase_     :     Any=None					)					->					Tuple:
  
  
  
  
  								"""simple docstring"""
  
  
  
  
  
  
  								if self.gradient_state.sync_gradients:
  													if self.accelerator_state.distributed_type == DistributedType.TPU:
  																		SCREAMING_SNAKE_CASE_	      =       {"closure": closure} if closure is not None else {}
  																		xm.optimizer_step(self.optimizer							,							optimizer_args=__A					)
  													elif self.scaler is not None:
  																		SCREAMING_SNAKE_CASE_	      =       False
  																		if self._last_scale is None:
  																							# `get_scale` is an async operation requiring full synchronization
  																							# on CPU and GPUs before finishing. As a result, we store away
  																							# the prior one to reduce the call overhead
  																							SCREAMING_SNAKE_CASE_	      =       self.scaler.get_scale()
  																							SCREAMING_SNAKE_CASE_	      =       True
  																		self.scaler.step(self.optimizer							,							__A					)
  																		self.scaler.update()
  																		SCREAMING_SNAKE_CASE_	      =       self.scaler.get_scale()
  																		if not new_scale:
  																							# If we reduced the loss scale, it means the optimizer step was skipped because of gradient overflow.
  																							SCREAMING_SNAKE_CASE_	      =       scale_after < self._last_scale
  																		SCREAMING_SNAKE_CASE_	      =       scale_after
  
  													else:
  																		self.optimizer.step(__A					)
  
  			def _lowercase (	self     :     Optional[Any]							,							lowerCAmelCase_     :     Any					)					->					List[str]:
  
  
  
  
  								"""simple docstring"""
  
  
  
  
  
  
  								for param_group in self.optimizer.param_groups:
  													SCREAMING_SNAKE_CASE_	      =       [parameters_map.get(__A							,							__A					) for p in param_group["params"]]
  
  			@property
  			def _lowercase (	self     :     Union[str, Any]					)					->					int:
  
  
  
  
  								"""simple docstring"""
  
  
  
  
  
  
  								warnings.warn(
  								    '''The `is_overflow` property is deprecated and will be removed in version 1.0 of Accelerate use '''
  								    '''`optimizer.step_was_skipped` instead.'''							,							__A							,							)
  								return self._is_overflow
  
  			@property
  			def _lowercase (	self     :     Tuple					)					->					str:
  
  
  
  
  								"""simple docstring"""
  
  
  
  
  
  
  								return self._is_overflow
  
  			def __getstate__(	self     :     Any					)					->					Optional[int]:
  
  
  
  
  								"""simple docstring"""
  
  
  
  
  
  
  								return self.__dict__.copy()
  
  
  
  
  			def __setstate__(	self     :     List[Any]							,							lowerCAmelCase_     :     Optional[int]					)					->					Any:
  
  
  
  
  								"""simple docstring"""
  
  
  
  
  
  
  								self.__dict__.update(__A					)
  
  
  
style_context: 
  
  
  def 							UpperCAmelCase				(     UpperCAmelCase				,UpperCAmelCase      )->		int:
  
  
  
  
  
  
  
  					'''simple docstring'''
  
  
  
  					return (input_a, input_a).count(0      ) % 2
  
  
  
  
  
  
  
  def 							UpperCAmelCase				(     )->		None:
  
  
  
  
  
  
  
  					'''simple docstring'''
  
  
  
  					assert xor_gate(0				,0      ) == 0
  					assert xor_gate(0				,1      ) == 1
  					assert xor_gate(1				,0      ) == 1
  					assert xor_gate(1				,1      ) == 0
  
  
  if __name__ == "__main__":
  		print(xor_gate(0, 0))
  		print(xor_gate(0, 1))
  
  
  