code: 
  import os
  import random
  from pathlib import Path
  from typing import List
  
  import numpy as np
  import torch
  from torch.cuda.amp import GradScaler
  
  from .utils import (
      MODEL_NAME,
      OPTIMIZER_NAME,
      RNG_STATE_NAME,
      SCALER_NAME,
      SCHEDULER_NAME,
      get_pretty_name,
      is_tpu_available,
      is_xpu_available,
      save,
  )
  
  
  if is_tpu_available(check_device=False):
  						import torch_xla.core.xla_model as xm
  
  from .logging import get_logger
  from .state import PartialState
  
  
  _snake_case			:       str      =						get_logger(__name__)
  
  def    __snake_case       (  __magic_name__		,						__magic_name__		,						__magic_name__		,						__magic_name__		,						__magic_name__		,						__magic_name__ = None		,						):
  
  
  
  
  
  
  
  				'''simple docstring'''
  
  
  				for i, state in enumerate(__magic_name__	):
  								lowercase						    =			F'''{MODEL_NAME}.bin''' if i == 0 else F'''{MODEL_NAME}_{i}.bin'''
  								lowercase						    =			os.path.join(__magic_name__		,						__magic_name__	)
  								save(__magic_name__		,						__magic_name__	)
  								logger.info(F'''Model weights saved in {output_model_file}'''	)
  				# Optimizer states
  				for i, opt in enumerate(__magic_name__	):
  								lowercase						    =			opt.state_dict()
  								lowercase						    =			F'''{OPTIMIZER_NAME}.bin''' if i == 0 else F'''{OPTIMIZER_NAME}_{i}.bin'''
  								lowercase						    =			os.path.join(__magic_name__		,						__magic_name__	)
  								save(__magic_name__		,						__magic_name__	)
  								logger.info(F'''Optimizer state saved in {output_optimizer_file}'''	)
  				# Scheduler states
  				for i, scheduler in enumerate(__magic_name__	):
  								lowercase						    =			scheduler.state_dict()
  								lowercase						    =			F'''{SCHEDULER_NAME}.bin''' if i == 0 else F'''{SCHEDULER_NAME}_{i}.bin'''
  								lowercase						    =			os.path.join(__magic_name__		,						__magic_name__	)
  								save(__magic_name__		,						__magic_name__	)
  								logger.info(F'''Scheduler state saved in {output_scheduler_file}'''	)
  				# GradScaler state
  				if scaler is not None:
  								lowercase						    =			scaler.state_dict()
  								lowercase						    =			os.path.join(__magic_name__		,						__magic_name__	)
  								torch.save(__magic_name__		,						__magic_name__	)
  								logger.info(F'''Gradient scaler state saved in {output_scaler_file}'''	)
  				# Random number generator states
  				lowercase						    =			{}
  				lowercase						    =			F'''{RNG_STATE_NAME}_{process_index}.pkl'''
  				lowercase						    =			random.getstate()
  				lowercase						    =			np.random.get_state()
  				lowercase						    =			torch.get_rng_state()
  				if is_xpu_available():
  								lowercase						    =			torch.xpu.get_rng_state_all()
  				else:
  								lowercase						    =			torch.cuda.get_rng_state_all()
  				if is_tpu_available():
  								lowercase						    =			xm.get_rng_state()
  				lowercase						    =			os.path.join(__magic_name__		,						__magic_name__	)
  				torch.save(__magic_name__		,						__magic_name__	)
  				logger.info(F'''Random states saved in {output_states_file}'''	)
  				return output_dir
  
  def    __snake_case       (  __magic_name__		,						__magic_name__		,						__magic_name__		,						__magic_name__		,						__magic_name__		,						__magic_name__=None		,						__magic_name__=None		,						**__magic_name__		,						):
  
  
  
  
  
  
  
  				'''simple docstring'''
  
  
  				if map_location not in [None, "cpu", "on_device"]:
  								raise TypeError(
  								    "Unsupported optimizer map location passed, please choose one of `None`, `'cpu'`, or `'on_device'`"	)
  				if map_location is None:
  								lowercase						    =			"cpu"
  				elif map_location == "on_device":
  								lowercase						    =			PartialState().device
  				# Model states
  				for i, model in enumerate(__magic_name__	):
  								lowercase						    =			F'''{MODEL_NAME}.bin''' if i == 0 else F'''{MODEL_NAME}_{i}.bin'''
  								lowercase						    =			os.path.join(__magic_name__		,						__magic_name__	)
  								models[i].load_state_dict(torch.load(__magic_name__		,						map_location=__magic_name__	)		,						**__magic_name__	)
  				logger.info("All model weights loaded successfully"	)
  
  				# Optimizer states
  				for i, opt in enumerate(__magic_name__	):
  								lowercase						    =			F'''{OPTIMIZER_NAME}.bin''' if i == 0 else F'''{OPTIMIZER_NAME}_{i}.bin'''
  								lowercase						    =			os.path.join(__magic_name__		,						__magic_name__	)
  								lowercase						    =			torch.load(__magic_name__		,						map_location=__magic_name__	)
  								optimizers[i].load_state_dict(__magic_name__	)
  				logger.info("All optimizer states loaded successfully"	)
  
  				# Scheduler states
  				for i, scheduler in enumerate(__magic_name__	):
  								lowercase						    =			F'''{SCHEDULER_NAME}.bin''' if i == 0 else F'''{SCHEDULER_NAME}_{i}.bin'''
  								lowercase						    =			os.path.join(__magic_name__		,						__magic_name__	)
  								scheduler.load_state_dict(torch.load(__magic_name__	)	)
  				logger.info("All scheduler states loaded successfully"	)
  
  				# GradScaler state
  				if scaler is not None:
  								lowercase						    =			os.path.join(__magic_name__		,						__magic_name__	)
  								scaler.load_state_dict(torch.load(__magic_name__	)	)
  								logger.info("GradScaler state loaded successfully"	)
  
  				# Random states
  				try:
  								lowercase						    =			torch.load(os.path.join(__magic_name__		,						F'''{RNG_STATE_NAME}_{process_index}.pkl'''	)	)
  								random.setstate(states["random_state"]	)
  								np.random.set_state(states["numpy_random_seed"]	)
  								torch.set_rng_state(states["torch_manual_seed"]	)
  								if is_xpu_available():
  												torch.xpu.set_rng_state_all(states["torch_xpu_manual_seed"]	)
  								else:
  												torch.cuda.set_rng_state_all(states["torch_cuda_manual_seed"]	)
  								if is_tpu_available():
  												xm.set_rng_state(states["xm_seed"]	)
  								logger.info("All random states loaded successfully"	)
  				except Exception:
  								logger.info("Could not load random states"	)
  
  def    __snake_case       (  __magic_name__		,						__magic_name__		,						__magic_name__ = 0	):
  
  
  
  
  
  
  
  				'''simple docstring'''
  
  
  				lowercase						    =			Path(__magic_name__	) / F'''custom_checkpoint_{index}.pkl'''
  				logger.info(F'''Saving the state of {get_pretty_name(__magic_name__	)} to {save_location}'''	)
  				torch.save(obj.state_dict()		,						__magic_name__	)
  
  
  
  
  def    __snake_case       (  __magic_name__		,						__magic_name__		,						__magic_name__ = 0	):
  
  
  
  
  
  
  
  				'''simple docstring'''
  
  
  				lowercase						    =			F'''{path}/custom_checkpoint_{index}.pkl'''
  				logger.info(F'''Loading the state of {get_pretty_name(__magic_name__	)} from {load_location}'''	)
  				obj.load_state_dict(torch.load(__magic_name__		,						map_location="cpu"	)	)
  
style_context: 
  from ...configuration_utils import PretrainedConfig
  from ...utils import logging
  
  
  _snake_case			:       List[Any]      =						logging.get_logger(__name__)
  
  _snake_case			:       Dict      =						{
      "google/reformer-crime-and-punishment": (
          "https://huggingface.co/google/reformer-crime-and-punishment/resolve/main/config.json"
      ),
      "google/reformer-enwik8": "https://huggingface.co/google/reformer-enwik8/resolve/main/config.json",
  }
  
  
  
  
  
  class      UpperCamelCase_			(			__a		):
  
  
  
  
  
  
  
  						'''simple docstring'''
  
  
  
  						UpperCamelCase			:     Any							      =				'''reformer'''
  						UpperCamelCase			:     List[str]							      =				['''past_buckets_states''']
  						UpperCamelCase			:     Optional[int]							      =				{}
  
  
  
  
  
  						def __init__( self       :Dict		,						lowerCAmelCase__       :Any=64		,						lowerCAmelCase__       :str=["local", "lsh", "local", "lsh", "local", "lsh"]		,						lowerCAmelCase__       :Optional[Any]=1.0		,						lowerCAmelCase__       :List[str]=True		,						lowerCAmelCase__       :Optional[Any]=[64, 64]		,						lowerCAmelCase__       :Tuple=[64, 192]		,						lowerCAmelCase__       :Union[str, Any]=0		,						lowerCAmelCase__       :Optional[Any]=2		,						lowerCAmelCase__       :List[str]=512		,						lowerCAmelCase__       :Any=None		,						lowerCAmelCase__       :int="relu"		,						lowerCAmelCase__       :Tuple=0.05		,						lowerCAmelCase__       :int=256		,						lowerCAmelCase__       :Optional[Any]=0.02		,						lowerCAmelCase__       :Optional[int]=False		,						lowerCAmelCase__       :Optional[Any]=1E-1_2		,						lowerCAmelCase__       :int=1		,						lowerCAmelCase__       :Optional[Any]=0		,						lowerCAmelCase__       :Any=0.05		,						lowerCAmelCase__       :Optional[int]=64		,						lowerCAmelCase__       :List[Any]=64		,						lowerCAmelCase__       :Tuple=0.0		,						lowerCAmelCase__       :Tuple=1		,						lowerCAmelCase__       :List[str]=0		,						lowerCAmelCase__       :Dict=4096		,						lowerCAmelCase__       :Tuple=12		,						lowerCAmelCase__       :Tuple=None		,						lowerCAmelCase__       :Union[str, Any]=1		,						lowerCAmelCase__       :int=0		,						lowerCAmelCase__       :Optional[Any]=320		,						lowerCAmelCase__       :Optional[Any]=False		,						lowerCAmelCase__       :Optional[int]=True		,						lowerCAmelCase__       :str=None		,						**lowerCAmelCase__       :Any		,						)     ->Any:
  										lowercase						    =			hash_seed
  										lowercase						    =			vocab_size
  										lowercase						    =			attention_head_size
  										lowercase						    =			hidden_size
  										lowercase						    =			num_attention_heads
  										lowercase						    =			num_hashes
  										lowercase						    =			len(lowerCAmelCase__    )
  										lowercase						    =			tuple(lowerCAmelCase__    ) if isinstance(lowerCAmelCase__		,						lowerCAmelCase__    ) else num_buckets
  										lowercase						    =			lsh_attn_chunk_length
  										lowercase						    =			local_attn_chunk_length
  										lowercase						    =			lsh_num_chunks_after
  										lowercase						    =			lsh_num_chunks_before
  										lowercase						    =			local_num_chunks_after
  										lowercase						    =			local_num_chunks_before
  										lowercase						    =			hidden_act
  										lowercase						    =			feed_forward_size
  										lowercase						    =			hidden_dropout_prob
  										lowercase						    =			lsh_attention_probs_dropout_prob
  										lowercase						    =			local_attention_probs_dropout_prob
  										lowercase						    =			max_position_embeddings
  										lowercase						    =			initializer_range
  										lowercase						    =			layer_norm_eps
  										lowercase						    =			axial_pos_embds
  										lowercase						    =			tuple(lowerCAmelCase__    )
  										lowercase						    =			tuple(lowerCAmelCase__    )
  										lowercase						    =			axial_norm_std
  										lowercase						    =			chunk_size_lm_head
  										lowercase						    =			attn_layers
  										lowercase						    =			use_cache
  										lowercase						    =			classifier_dropout
  										super().__init__(
  										    pad_token_id=lowerCAmelCase__		,						eos_token_id=lowerCAmelCase__		,						is_decoder=lowerCAmelCase__		,						tie_word_embeddings=lowerCAmelCase__		,						**lowerCAmelCase__		,						)
  