code: 
  
  
  
  """simple docstring"""
  import unittest
  from functools import partial
  
  import torch
  
  from accelerate import Accelerator, debug_launcher
  from accelerate.state import AcceleratorState, GradientState
  from accelerate.test_utils import require_cpu, require_huggingface_suite
  from accelerate.utils import GradientAccumulationPlugin
  
  
  
  
  
  def 				lowerCamelCase_		(UpperCamelCase__	:	Dict=2  , UpperCamelCase__	:	Any=True  , UpperCamelCase__	:	Optional[int]=False							):
      _UpperCAmelCase							:						Dict       							=					Accelerator(step_scheduler_with_optimizer=UpperCamelCase__  , split_batches=UpperCamelCase__							)
      _UpperCAmelCase							:						str       							=					torch.nn.Linear(2  , 4							)
      _UpperCAmelCase							:						Tuple       							=					torch.optim.AdamW(model.parameters()  , lr=1.0							)
      _UpperCAmelCase							:						List[str]       							=					torch.optim.lr_scheduler.OneCycleLR(UpperCamelCase__  , max_lr=0.01  , steps_per_epoch=2  , epochs=1							)
      _UpperCAmelCase							,       _UpperCAmelCase							,       _UpperCAmelCase							:						Any       							=					accelerator.prepare(UpperCamelCase__  , UpperCamelCase__  , UpperCamelCase__							)
  
      # Optimizer has stepped
      scheduler.step()
      if step_scheduler_with_optimizer or (num_processes == 1):
          assert (
              scheduler.scheduler.last_epoch == num_processes
          ), F'Last Epoch ({scheduler.scheduler.last_epoch}) != Num Processes ({num_processes})'
      else:
          assert (
              scheduler.scheduler.last_epoch != num_processes
          ), F'Last Epoch ({scheduler.scheduler.last_epoch}) == Num Processes ({num_processes})'
  
  
  
  
  
  def 				lowerCamelCase_		(UpperCamelCase__	:	Optional[int]=2  , UpperCamelCase__	:	List[str]=True  , UpperCamelCase__	:	Optional[Any]=False							):
      _UpperCAmelCase							:						Any       							=					Accelerator(step_scheduler_with_optimizer=UpperCamelCase__  , split_batches=UpperCamelCase__							)
      _UpperCAmelCase							:						Tuple       							=					torch.nn.Linear(2  , 4							)
      _UpperCAmelCase							:						Tuple       							=					torch.optim.AdamW(model.parameters()  , lr=1.0							)
      _UpperCAmelCase							:						List[str]       							=					torch.optim.lr_scheduler.LambdaLR(UpperCamelCase__  , lr_lambda=lambda UpperCamelCase__							: 1 - n / 10							)
      _UpperCAmelCase							,       _UpperCAmelCase							,       _UpperCAmelCase							:						Tuple       							=					accelerator.prepare(UpperCamelCase__  , UpperCamelCase__  , UpperCamelCase__							)
  
      # Optimizer has stepped
      _UpperCAmelCase							:						List[str]       							=					False
      scheduler.step()
      _UpperCAmelCase							:						List[str]       							=					1 - (num_processes if (step_scheduler_with_optimizer and not split_batches) else 1) / 10
      assert (
          scheduler.get_last_lr()[0] == expected_lr
      ), F'Wrong lr found at first step, expected {expected_lr}, got {scheduler.get_last_lr()[0]}'
  
      # Optimizer has not stepped
      _UpperCAmelCase							:						int       							=					True
      scheduler.step()
      if not step_scheduler_with_optimizer:
          _UpperCAmelCase							:						Tuple       							=					1 - 2 / 10
      assert (
          scheduler.get_last_lr()[0] == expected_lr
      ), F'Wrong lr found at second step, expected {expected_lr}, got {scheduler.get_last_lr()[0]}'
  
  
  
  
  
  def 				lowerCamelCase_		(UpperCamelCase__	:	int = 2							):
      from transformers import get_linear_schedule_with_warmup
  
      _UpperCAmelCase							:						List[str]       							=					[1, 2, 4]
      for num_steps in steps:
          _UpperCAmelCase							:						Optional[int]       							=					GradientAccumulationPlugin(num_steps=UpperCamelCase__  , adjust_scheduler=num_steps > 1							)
          _UpperCAmelCase							:						List[Any]       							=					Accelerator(gradient_accumulation_plugin=UpperCamelCase__							)
          _UpperCAmelCase							:						Dict       							=					torch.nn.Linear(2  , 4							)
          _UpperCAmelCase							:						Any       							=					torch.optim.AdamW(model.parameters()  , lr=10.0							)
          _UpperCAmelCase							:						Optional[Any]       							=					get_linear_schedule_with_warmup(optimizer=UpperCamelCase__  , num_warmup_steps=0  , num_training_steps=20							)
  
          _UpperCAmelCase							,       _UpperCAmelCase							,       _UpperCAmelCase							:						Optional[Any]       							=					accelerator.prepare(UpperCamelCase__  , UpperCamelCase__  , UpperCamelCase__							)
  
          for i in range(10 * num_steps							):
              with accelerator.accumulate(UpperCamelCase__							):
                  optimizer.step()
                  scheduler.step()
  
              if i == (10 * num_steps - 2):
                  assert (
                      scheduler.get_last_lr()[0] != 0
                  ), F'Wrong lr found at second-to-last step, expected non-zero, got {scheduler.get_last_lr()[0]}. num_steps: {num_steps}'
          assert (
              scheduler.get_last_lr()[0] == 0
          ), F'Wrong lr found at last step, expected 0, got {scheduler.get_last_lr()[0]}'
          GradientState._reset_state()
  
  
  @require_cpu
  class 							_UpperCAmelCase							(				unittest.TestCase   ):
  
  
  
  
  
       '''simple docstring'''
  
  
  
       def 				__lowerCAmelCase			(      self	)							->      Optional[Any]:
           debug_launcher(partial(A ,			num_processes=1	) ,			num_processes=1	)
           debug_launcher(partial(A ,			num_processes=1 ,			split_batches=A	) ,			num_processes=1	)
  
  
  
  
  
       def 				__lowerCAmelCase			(      self	)							->      List[Any]:
           debug_launcher(partial(A ,			num_processes=1	) ,			num_processes=1	)
           debug_launcher(partial(A ,			num_processes=1 ,			split_batches=A	) ,			num_processes=1	)
  
  
  
  
  
       def 				__lowerCAmelCase			(      self	)							->      str:
           debug_launcher(partial(A ,			num_processes=1 ,			step_scheduler_with_optimizer=A	) ,			num_processes=1	)
  
  
  
  
  
       def 				__lowerCAmelCase			(      self	)							->      Optional[int]:
           debug_launcher(partial(A ,			num_processes=1 ,			step_scheduler_with_optimizer=A	) ,			num_processes=1	)
  
  
  
  
  
       def 				__lowerCAmelCase			(      self	)							->      List[str]:
           AcceleratorState._reset_state(A	)
           debug_launcher(A	)
           debug_launcher(partial(A ,			num_processes=1 ,			split_batches=A	) ,			num_processes=1	)
  
  
  
  
  
       def 				__lowerCAmelCase			(      self	)							->      List[str]:
           AcceleratorState._reset_state(A	)
           debug_launcher(A	)
           debug_launcher(partial(A ,			num_processes=1 ,			split_batches=A	) ,			num_processes=1	)
  
  
  
  
  
       def 				__lowerCAmelCase			(      self	)							->      Optional[Any]:
           AcceleratorState._reset_state(A	)
           debug_launcher(partial(A ,			step_scheduler_with_optimizer=A	)	)
  
  
  
  
  
       def 				__lowerCAmelCase			(      self	)							->      Optional[Any]:
           AcceleratorState._reset_state(A	)
           debug_launcher(partial(A ,			step_scheduler_with_optimizer=A	)	)
  
  
  
  
  
       @require_huggingface_suite
       def 				__lowerCAmelCase			(      self	)							->      Dict:
           AcceleratorState._reset_state(A	)
           debug_launcher(partial(A ,			num_processes=1	)	)
           debug_launcher(A	)
  
  
  
style_context: 
  
  
  
  """simple docstring"""
  import logging
  
  import torch
  from torch import nn
  from torch.nn import CrossEntropyLoss, MSELoss
  
  from transformers.file_utils import add_start_docstrings, add_start_docstrings_to_model_forward
  from transformers.models.albert.modeling_albert import (
      ALBERT_INPUTS_DOCSTRING,
      ALBERT_START_DOCSTRING,
      AlbertModel,
      AlbertPreTrainedModel,
      AlbertTransformer,
  )
  
  
  _lowerCAmelCase					:Union[str, Any]  =      logging.getLogger(__name__)
  class 							_UpperCAmelCase							(				a   ):
  
  
  
  
  
       '''simple docstring'''
  
  
  
       def 				__lowerCAmelCase			(      self ,			A ,			A ,			A=None ,			A=None	)							->      Tuple:
           if current_layer == 0:
               _UpperCAmelCase							:						Optional[Any]       							=					self.embedding_hidden_mapping_in(A	)
           else:
               _UpperCAmelCase							:						Union[str, Any]       							=					hidden_states[0]
  
           _UpperCAmelCase							:						Dict       							=					int(self.config.num_hidden_layers / self.config.num_hidden_groups	)
  
           # Index of the hidden group
           _UpperCAmelCase							:						Dict       							=					int(current_layer / (self.config.num_hidden_layers / self.config.num_hidden_groups)	)
  
           _UpperCAmelCase							:						List[Any]       							=					self.albert_layer_groups[group_idx](
               A ,			A ,			head_mask[group_idx * layers_per_group : (group_idx + 1) * layers_per_group] ,			)
           _UpperCAmelCase							:						str       							=					layer_group_output[0]
  
           return (hidden_states,)
  
  
  
  @add_start_docstrings(
      '''The bare ALBERT Model transformer with PABEE outputting raw hidden-states without any specific head on top.'''	,a	,)
  class 							_UpperCAmelCase							(				a   ):
  
  
  
  
  
       '''simple docstring'''
  
  
  
       def __init__(      self ,			A	)							->      str:
           super().__init__(A	)
  
           _UpperCAmelCase							:						Dict       							=					AlbertTransformerWithPabee(A	)
  
           self.init_weights()
           _UpperCAmelCase							:						str       							=					0
           _UpperCAmelCase							:						Union[str, Any]       							=					0
           _UpperCAmelCase							:						Dict       							=					0
  
           _UpperCAmelCase							:						str       							=					0
  
  
  
  
  
       def 				__lowerCAmelCase			(      self ,			A	)							->      Optional[int]:
           _UpperCAmelCase							:						List[Any]       							=					threshold
  
  
  
  
  
       def 				__lowerCAmelCase			(      self ,			A	)							->      Optional[Any]:
           _UpperCAmelCase							:						int       							=					patience
  
  
  
  
  
       def 				__lowerCAmelCase			(      self	)							->      Any:
           _UpperCAmelCase							:						Dict       							=					0
           _UpperCAmelCase							:						List[str]       							=					0
  
  
  
  
  
       def 				__lowerCAmelCase			(      self	)							->      List[Any]:
           _UpperCAmelCase							:						Union[str, Any]       							=					self.inference_layers_num / self.inference_instances_num
           _UpperCAmelCase							:						str       							=					(
               f'*** Patience = {self.patience} Avg. Inference Layers = {avg_inf_layers:.2f} Speed Up ='
               f' {1 - avg_inf_layers / self.config.num_hidden_layers:.2f} ***'
           )
           print(A	)
  
  
  
  
  
       @add_start_docstrings_to_model_forward(A	)
       def 				__lowerCAmelCase			(      self ,			A=None ,			A=None ,			A=None ,			A=None ,			A=None ,			A=None ,			A=None ,			A=None ,			A=False ,			)							->      int:
  
           if input_ids is not None and inputs_embeds is not None:
               raise ValueError('''You cannot specify both input_ids and inputs_embeds at the same time'''	)
           elif input_ids is not None:
               _UpperCAmelCase							:						Optional[Any]       							=					input_ids.size()
           elif inputs_embeds is not None:
               _UpperCAmelCase							:						Optional[Any]       							=					inputs_embeds.size()[:-1]
           else:
               raise ValueError('''You have to specify either input_ids or inputs_embeds'''	)
  
           _UpperCAmelCase							:						Tuple       							=					input_ids.device if input_ids is not None else inputs_embeds.device
  
           if attention_mask is None:
               _UpperCAmelCase							:						Union[str, Any]       							=					torch.ones(A ,			device=A	)
           if token_type_ids is None:
               _UpperCAmelCase							:						Tuple       							=					torch.zeros(A ,			dtype=torch.long ,			device=A	)
  
           _UpperCAmelCase							:						Optional[Any]       							=					attention_mask.unsqueeze(1	).unsqueeze(2	)
           _UpperCAmelCase							:						Union[str, Any]       							=					extended_attention_mask.to(dtype=self.dtype	)  # fp16 compatibility
           _UpperCAmelCase							:						Dict       							=					(1.0 - extended_attention_mask) * -10_000.0
           _UpperCAmelCase							:						Any       							=					self.get_head_mask(A ,			self.config.num_hidden_layers	)
  
           _UpperCAmelCase							:						int       							=					self.embeddings(
               A ,			position_ids=A ,			token_type_ids=A ,			inputs_embeds=A	)
           _UpperCAmelCase							:						Any       							=					embedding_output
  
           if self.training:
               _UpperCAmelCase							:						Optional[Any]       							=					[]
               for i in range(self.config.num_hidden_layers	):
                   _UpperCAmelCase							:						Union[str, Any]       							=					self.encoder.adaptive_forward(
                       A ,			current_layer=A ,			attention_mask=A ,			head_mask=A ,			)
  
                   _UpperCAmelCase							:						List[str]       							=					self.pooler_activation(self.pooler(encoder_outputs[0][:, 0]	)	)
                   _UpperCAmelCase							:						Dict       							=					output_layers[i](output_dropout(A	)	)
                   res.append(A	)
           elif self.patience == 0:  # Use all layers for inference
               _UpperCAmelCase							:						Dict       							=					self.encoder(A ,			A ,			head_mask=A	)
               _UpperCAmelCase							:						Tuple       							=					self.pooler_activation(self.pooler(encoder_outputs[0][:, 0]	)	)
               _UpperCAmelCase							:						Optional[int]       							=					[output_layers[self.config.num_hidden_layers - 1](A	)]
           else:
               _UpperCAmelCase							:						List[Any]       							=					0
               _UpperCAmelCase							:						int       							=					None
               _UpperCAmelCase							:						int       							=					0
               for i in range(self.config.num_hidden_layers	):
                   calculated_layer_num += 1
                   _UpperCAmelCase							:						Tuple       							=					self.encoder.adaptive_forward(
                       A ,			current_layer=A ,			attention_mask=A ,			head_mask=A ,			)
  
                   _UpperCAmelCase							:						Any       							=					self.pooler_activation(self.pooler(encoder_outputs[0][:, 0]	)	)
                   _UpperCAmelCase							:						Optional[Any]       							=					output_layers[i](A	)
                   if regression:
                       _UpperCAmelCase							:						str       							=					logits.detach()
                       if patient_result is not None:
                           _UpperCAmelCase							:						str       							=					patient_result.detach()
                       if (patient_result is not None) and torch.abs(patient_result - labels	) < self.regression_threshold:
                           patient_counter += 1
                       else:
                           _UpperCAmelCase							:						List[Any]       							=					0
                   else:
                       _UpperCAmelCase							:						List[str]       							=					logits.detach().argmax(dim=1	)
                       if patient_result is not None:
                           _UpperCAmelCase							:						Optional[Any]       							=					patient_result.detach().argmax(dim=1	)
                       if (patient_result is not None) and torch.all(labels.eq(A	)	):
                           patient_counter += 1
                       else:
                           _UpperCAmelCase							:						Tuple       							=					0
  
                   _UpperCAmelCase							:						Any       							=					logits
                   if patient_counter == self.patience:
                       break
               _UpperCAmelCase							:						List[str]       							=					[patient_result]
               self.inference_layers_num += calculated_layer_num
               self.inference_instances_num += 1
  
           return res
  
  
  
  @add_start_docstrings(
      '''Albert Model transformer with PABEE and a sequence classification/regression head on top (a linear layer on top of
      the pooled output) e.g. for GLUE tasks. '''	,a	,)
  class 							_UpperCAmelCase							(				a   ):
  
  
  
  
  
       '''simple docstring'''
  
  
  
       def __init__(      self ,			A	)							->      Optional[int]:
           super().__init__(A	)
           _UpperCAmelCase							:						List[Any]       							=					config.num_labels
  
           _UpperCAmelCase							:						Optional[int]       							=					AlbertModelWithPabee(A	)
           _UpperCAmelCase							:						List[str]       							=					nn.Dropout(config.classifier_dropout_prob	)
           _UpperCAmelCase							:						Optional[Any]       							=					nn.ModuleList(
               [nn.Linear(config.hidden_size ,			self.config.num_labels	) for _ in range(config.num_hidden_layers	)]	)
  
           self.init_weights()
  
  
  
  
  
       @add_start_docstrings_to_model_forward(A	)
       def 				__lowerCAmelCase			(      self ,			A=None ,			A=None ,			A=None ,			A=None ,			A=None ,			A=None ,			A=None ,			)							->      List[str]:
           _UpperCAmelCase							:						Optional[Any]       							=					self.albert(
               input_ids=A ,			attention_mask=A ,			token_type_ids=A ,			position_ids=A ,			head_mask=A ,			inputs_embeds=A ,			output_dropout=self.dropout ,			output_layers=self.classifiers ,			regression=self.num_labels == 1 ,			)
  
           _UpperCAmelCase							:						int       							=					(logits[-1],)
  
           if labels is not None:
               _UpperCAmelCase							:						int       							=					None
               _UpperCAmelCase							:						Tuple       							=					0
               for ix, logits_item in enumerate(A	):
                   if self.num_labels == 1:
                       #  We are doing regression
                       _UpperCAmelCase							:						Dict       							=					MSELoss()
                       _UpperCAmelCase							:						Optional[int]       							=					loss_fct(logits_item.view(-1	) ,			labels.view(-1	)	)
                   else:
                       _UpperCAmelCase							:						List[str]       							=					CrossEntropyLoss()
                       _UpperCAmelCase							:						Any       							=					loss_fct(logits_item.view(-1 ,			self.num_labels	) ,			labels.view(-1	)	)
                   if total_loss is None:
                       _UpperCAmelCase							:						Optional[Any]       							=					loss
                   else:
                       total_loss += loss * (ix + 1)
                   total_weights += ix + 1
               _UpperCAmelCase							:						Any       							=					(total_loss / total_weights,) + outputs
  
           return outputs
  
  
  