code: 
  """simple docstring"""
  
  
  
  
  
  
  
  import argparse
  
  import torch
  
  from diffusers import UNetaDConditionModel
  def 			__UpperCAmelCase		(		a_			:		Optional[int]    ,  a_			:		str    ,  a_			:		List[str]    ,  a_			:		List[str]=None    ,  a_			:		Any=None    ,  a_			:		str=None				)     ->       int:
     assert isinstance(a_    ,  a_				), "Paths should be a list of dicts containing 'old' and 'new' keys."
  
     # Splits the attention layers into three variables.
     if attention_paths_to_split is not None:
        for path, path_map in attention_paths_to_split.items():
           _UpperCAmelCase   	=					old_checkpoint[path]
           _UpperCAmelCase   	=					old_tensor.shape[0] // 3
  
           _UpperCAmelCase   	=					(-1, channels) if len(old_tensor.shape				) == 3 else (-1)
  
           _UpperCAmelCase   	=					old_tensor.shape[0] // config["num_head_channels"] // 3
  
           _UpperCAmelCase   	=					old_tensor.reshape((num_heads, 3 * channels // num_heads) + old_tensor.shape[1:]				)
           _UpperCAmelCase	,							_UpperCAmelCase	,							_UpperCAmelCase   	=					old_tensor.split(channels // num_heads    ,  dim=1				)
  
           _UpperCAmelCase   	=					query.reshape(a_				)
           _UpperCAmelCase   	=					key.reshape(a_				)
           _UpperCAmelCase   	=					value.reshape(a_				)
  
     for path in paths:
        _UpperCAmelCase   	=					path["new"]
  
        # These have already been assigned
        if attention_paths_to_split is not None and new_path in attention_paths_to_split:
           continue
  
        if additional_replacements is not None:
           for replacement in additional_replacements:
              _UpperCAmelCase   	=					new_path.replace(replacement["old"]    ,  replacement["new"]				)
  
          # proj_attn.weight has to be converted from conv 1D to linear
        _UpperCAmelCase   	=					old_checkpoint[path["old"]]
        _UpperCAmelCase   	=					["proj_attn.weight"]
        _UpperCAmelCase   	=					["proj_out.weight", "proj_in.weight"]
        if any(k in new_path for k in names				):
           _UpperCAmelCase   	=					weight[:, :, 0]
        elif any(k in new_path for k in names_a				) and len(weight.shape				) > 2 and ".attentions." not in new_path:
           _UpperCAmelCase   	=					weight[:, :, 0]
        else:
           _UpperCAmelCase   	=					weight
  def 			__UpperCAmelCase		(		a_			:		Tuple    ,  a_			:		Any=0				)     ->       List[Any]:
     _UpperCAmelCase   	=					[]
     for old_item in old_list:
        _UpperCAmelCase   	=					old_item
  
        #         new_item = new_item.replace('norm.weight', 'group_norm.weight')
        #         new_item = new_item.replace('norm.bias', 'group_norm.bias')
  
        #         new_item = new_item.replace('proj_out.weight', 'proj_attn.weight')
        #         new_item = new_item.replace('proj_out.bias', 'proj_attn.bias')
  
        #         new_item = shave_segments(new_item, n_shave_prefix_segments=n_shave_prefix_segments)
  
        mapping.append({"old": old_item, "new": new_item}				)
  
     return mapping
  def 			__UpperCAmelCase		(		a_			:		Optional[int]    ,  a_			:		Dict=1				)     ->       List[str]:
     if n_shave_prefix_segments >= 0:
        return ".".join(path.split("."				)[n_shave_prefix_segments:]				)
     else:
        return ".".join(path.split("."				)[:n_shave_prefix_segments]				)
  def 			__UpperCAmelCase		(		a_			:		Union[str, Any]    ,  a_			:		int=0				)     ->       Optional[int]:
     _UpperCAmelCase   	=					[]
     for old_item in old_list:
        mapping.append({"old": old_item, "new": old_item}				)
  
     return mapping
  def 			__UpperCAmelCase		(		a_			:		List[Any]    ,  a_			:		Any=0				)     ->       Union[str, Any]:
     _UpperCAmelCase   	=					[]
     for old_item in old_list:
        _UpperCAmelCase   	=					old_item.replace("in_layers.0"    ,  "norm1"				)
        _UpperCAmelCase   	=					new_item.replace("in_layers.2"    ,  "conv1"				)
  
        _UpperCAmelCase   	=					new_item.replace("out_layers.0"    ,  "norm2"				)
        _UpperCAmelCase   	=					new_item.replace("out_layers.3"    ,  "conv2"				)
  
        _UpperCAmelCase   	=					new_item.replace("emb_layers.1"    ,  "time_emb_proj"				)
        _UpperCAmelCase   	=					new_item.replace("skip_connection"    ,  "conv_shortcut"				)
  
        _UpperCAmelCase   	=					shave_segments(a_    ,  n_shave_prefix_segments=a_				)
  
        if "temopral_conv" not in old_item:
           mapping.append({"old": old_item, "new": new_item}				)
  
     return mapping
  def 			__UpperCAmelCase		(		a_			:		str    ,  a_			:		str    ,  a_			:		str=None    ,  a_			:		Tuple=False				)     ->       Optional[Any]:
     _UpperCAmelCase   	=					{}
     _UpperCAmelCase   	=					list(checkpoint.keys()				)
  
     _UpperCAmelCase   	=					"model.diffusion_model."
  
     # at least a 100 parameters have to start with `model_ema` in order for the checkpoint to be EMA
     if sum(k.startswith("model_ema"				) for k in keys				) > 1_00 and extract_ema:
        print(F'''Checkpoint {path} has both EMA and non-EMA weights.'''				)
        print(
            "In this conversion only the EMA weights are extracted. If you want to instead extract the non-EMA"
            " weights (useful to continue fine-tuning), please make sure to remove the `--extract_ema` flag."				)
        for key in keys:
           if key.startswith("model.diffusion_model"				):
              _UpperCAmelCase   	=					"model_ema." + "".join(key.split("."				)[1:]				)
              _UpperCAmelCase   	=					checkpoint.pop(a_				)
     else:
        if sum(k.startswith("model_ema"				) for k in keys				) > 1_00:
           print(
               "In this conversion only the non-EMA weights are extracted. If you want to instead extract the EMA"
               " weights (usually better for inference), please make sure to add the `--extract_ema` flag."				)
  
        for key in keys:
           _UpperCAmelCase   	=					checkpoint.pop(a_				)
  
     _UpperCAmelCase   	=					{}
  
     _UpperCAmelCase   	=					unet_state_dict["time_embed.0.weight"]
     _UpperCAmelCase   	=					unet_state_dict["time_embed.0.bias"]
     _UpperCAmelCase   	=					unet_state_dict["time_embed.2.weight"]
     _UpperCAmelCase   	=					unet_state_dict["time_embed.2.bias"]
  
     if config["class_embed_type"] is None:
        # No parameters to port
        ...
     elif config["class_embed_type"] == "timestep" or config["class_embed_type"] == "projection":
        _UpperCAmelCase   	=					unet_state_dict["label_emb.0.0.weight"]
        _UpperCAmelCase   	=					unet_state_dict["label_emb.0.0.bias"]
        _UpperCAmelCase   	=					unet_state_dict["label_emb.0.2.weight"]
        _UpperCAmelCase   	=					unet_state_dict["label_emb.0.2.bias"]
     else:
        raise NotImplementedError(F'''Not implemented `class_embed_type`: {config['class_embed_type']}'''				)
  
     _UpperCAmelCase   	=					unet_state_dict["input_blocks.0.0.weight"]
     _UpperCAmelCase   	=					unet_state_dict["input_blocks.0.0.bias"]
  
     _UpperCAmelCase   	=					[v for v in unet_state_dict if v.startswith("input_blocks.0.1"				)]
     _UpperCAmelCase   	=					renew_attention_paths(a_				)
     _UpperCAmelCase   	=					{"old": "input_blocks.0.1", "new": "transformer_in"}
     assign_to_checkpoint(a_    ,  a_    ,  a_    ,  additional_replacements=[meta_path]    ,  config=a_				)
  
     _UpperCAmelCase   	=					unet_state_dict["out.0.weight"]
     _UpperCAmelCase   	=					unet_state_dict["out.0.bias"]
     _UpperCAmelCase   	=					unet_state_dict["out.2.weight"]
     _UpperCAmelCase   	=					unet_state_dict["out.2.bias"]
  
     # Retrieves the keys for the input blocks only
     _UpperCAmelCase   	=					len({".".join(layer.split("."				)[:2]				) for layer in unet_state_dict if "input_blocks" in layer}				)
     _UpperCAmelCase   	=					{
         layer_id: [key for key in unet_state_dict if F'''input_blocks.{layer_id}''' in key]
         for layer_id in range(a_				)
     }
  
     # Retrieves the keys for the middle blocks only
     _UpperCAmelCase   	=					len({".".join(layer.split("."				)[:2]				) for layer in unet_state_dict if "middle_block" in layer}				)
     _UpperCAmelCase   	=					{
         layer_id: [key for key in unet_state_dict if F'''middle_block.{layer_id}''' in key]
         for layer_id in range(a_				)
     }
  
     # Retrieves the keys for the output blocks only
     _UpperCAmelCase   	=					len({".".join(layer.split("."				)[:2]				) for layer in unet_state_dict if "output_blocks" in layer}				)
     _UpperCAmelCase   	=					{
         layer_id: [key for key in unet_state_dict if F'''output_blocks.{layer_id}''' in key]
         for layer_id in range(a_				)
     }
  
     for i in range(1    ,  a_				):
        _UpperCAmelCase   	=					(i - 1) // (config["layers_per_block"] + 1)
        _UpperCAmelCase   	=					(i - 1) % (config["layers_per_block"] + 1)
  
        _UpperCAmelCase   	=					[
            key for key in input_blocks[i] if F'''input_blocks.{i}.0''' in key and F'''input_blocks.{i}.0.op''' not in key
        ]
        _UpperCAmelCase   	=					[key for key in input_blocks[i] if F'''input_blocks.{i}.1''' in key]
        _UpperCAmelCase   	=					[key for key in input_blocks[i] if F'''input_blocks.{i}.2''' in key]
  
        if F'''input_blocks.{i}.op.weight''' in unet_state_dict:
           _UpperCAmelCase   	=					unet_state_dict.pop(
               F'''input_blocks.{i}.op.weight'''				)
           _UpperCAmelCase   	=					unet_state_dict.pop(
               F'''input_blocks.{i}.op.bias'''				)
  
        _UpperCAmelCase   	=					renew_resnet_paths(a_				)
        _UpperCAmelCase   	=					{"old": F'''input_blocks.{i}.0''', "new": F'''down_blocks.{block_id}.resnets.{layer_in_block_id}'''}
        assign_to_checkpoint(
            a_    ,  a_    ,  a_    ,  additional_replacements=[meta_path]    ,  config=a_				)
  
        _UpperCAmelCase   	=					[key for key in resnets if "temopral_conv" in key]
        _UpperCAmelCase   	=					renew_temp_conv_paths(a_				)
        _UpperCAmelCase   	=					{
            "old": F'''input_blocks.{i}.0.temopral_conv''',
            "new": F'''down_blocks.{block_id}.temp_convs.{layer_in_block_id}''',
        }
        assign_to_checkpoint(
            a_    ,  a_    ,  a_    ,  additional_replacements=[meta_path]    ,  config=a_				)
  
        if len(a_				):
           _UpperCAmelCase   	=					renew_attention_paths(a_				)
           _UpperCAmelCase   	=					{"old": F'''input_blocks.{i}.1''', "new": F'''down_blocks.{block_id}.attentions.{layer_in_block_id}'''}
           assign_to_checkpoint(
               a_    ,  a_    ,  a_    ,  additional_replacements=[meta_path]    ,  config=a_				)
  
        if len(a_				):
           _UpperCAmelCase   	=					renew_attention_paths(a_				)
           _UpperCAmelCase   	=					{
               "old": F'''input_blocks.{i}.2''',
               "new": F'''down_blocks.{block_id}.temp_attentions.{layer_in_block_id}''',
           }
           assign_to_checkpoint(
               a_    ,  a_    ,  a_    ,  additional_replacements=[meta_path]    ,  config=a_				)
  
     _UpperCAmelCase   	=					middle_blocks[0]
     _UpperCAmelCase   	=					[key for key in resnet_a if "temopral_conv" in key]
     _UpperCAmelCase   	=					middle_blocks[1]
     _UpperCAmelCase   	=					middle_blocks[2]
     _UpperCAmelCase   	=					middle_blocks[3]
     _UpperCAmelCase   	=					[key for key in resnet_a if "temopral_conv" in key]
  
     _UpperCAmelCase   	=					renew_resnet_paths(a_				)
     _UpperCAmelCase   	=					{"old": "middle_block.0", "new": "mid_block.resnets.0"}
     assign_to_checkpoint(
         a_    ,  a_    ,  a_    ,  config=a_    ,  additional_replacements=[meta_path]				)
  
     _UpperCAmelCase   	=					renew_temp_conv_paths(a_				)
     _UpperCAmelCase   	=					{"old": "middle_block.0.temopral_conv", "new": "mid_block.temp_convs.0"}
     assign_to_checkpoint(
         a_    ,  a_    ,  a_    ,  config=a_    ,  additional_replacements=[meta_path]				)
  
     _UpperCAmelCase   	=					renew_resnet_paths(a_				)
     _UpperCAmelCase   	=					{"old": "middle_block.3", "new": "mid_block.resnets.1"}
     assign_to_checkpoint(
         a_    ,  a_    ,  a_    ,  config=a_    ,  additional_replacements=[meta_path]				)
  
     _UpperCAmelCase   	=					renew_temp_conv_paths(a_				)
     _UpperCAmelCase   	=					{"old": "middle_block.3.temopral_conv", "new": "mid_block.temp_convs.1"}
     assign_to_checkpoint(
         a_    ,  a_    ,  a_    ,  config=a_    ,  additional_replacements=[meta_path]				)
  
     _UpperCAmelCase   	=					renew_attention_paths(a_				)
     _UpperCAmelCase   	=					{"old": "middle_block.1", "new": "mid_block.attentions.0"}
     assign_to_checkpoint(
         a_    ,  a_    ,  a_    ,  additional_replacements=[meta_path]    ,  config=a_				)
  
     _UpperCAmelCase   	=					renew_attention_paths(a_				)
     _UpperCAmelCase   	=					{"old": "middle_block.2", "new": "mid_block.temp_attentions.0"}
     assign_to_checkpoint(
         a_    ,  a_    ,  a_    ,  additional_replacements=[meta_path]    ,  config=a_				)
  
     for i in range(a_				):
        _UpperCAmelCase   	=					i // (config["layers_per_block"] + 1)
        _UpperCAmelCase   	=					i % (config["layers_per_block"] + 1)
        _UpperCAmelCase   	=					[shave_segments(a_    ,  2				) for name in output_blocks[i]]
        _UpperCAmelCase   	=					{}
  
        for layer in output_block_layers:
           _UpperCAmelCase	,							_UpperCAmelCase   	=					layer.split("."				)[0], shave_segments(a_    ,  1				)
           if layer_id in output_block_list:
              output_block_list[layer_id].append(a_				)
           else:
              _UpperCAmelCase   	=					[layer_name]
  
        if len(a_				) > 1:
           _UpperCAmelCase   	=					[key for key in output_blocks[i] if F'''output_blocks.{i}.0''' in key]
           _UpperCAmelCase   	=					[key for key in output_blocks[i] if F'''output_blocks.{i}.1''' in key]
           _UpperCAmelCase   	=					[key for key in output_blocks[i] if F'''output_blocks.{i}.2''' in key]
  
           _UpperCAmelCase   	=					renew_resnet_paths(a_				)
           _UpperCAmelCase   	=					renew_resnet_paths(a_				)
  
           _UpperCAmelCase   	=					{"old": F'''output_blocks.{i}.0''', "new": F'''up_blocks.{block_id}.resnets.{layer_in_block_id}'''}
           assign_to_checkpoint(
               a_    ,  a_    ,  a_    ,  additional_replacements=[meta_path]    ,  config=a_				)
  
           _UpperCAmelCase   	=					[key for key in resnets if "temopral_conv" in key]
           _UpperCAmelCase   	=					renew_temp_conv_paths(a_				)
           _UpperCAmelCase   	=					{
               "old": F'''output_blocks.{i}.0.temopral_conv''',
               "new": F'''up_blocks.{block_id}.temp_convs.{layer_in_block_id}''',
           }
           assign_to_checkpoint(
               a_    ,  a_    ,  a_    ,  additional_replacements=[meta_path]    ,  config=a_				)
  
           _UpperCAmelCase   	=					{k: sorted(a_				) for k, v in output_block_list.items()}
           if ["conv.bias", "conv.weight"] in output_block_list.values():
              _UpperCAmelCase   	=					list(output_block_list.values()				).index(["conv.bias", "conv.weight"]				)
              _UpperCAmelCase   	=					unet_state_dict[
                  F'''output_blocks.{i}.{index}.conv.weight'''
              ]
              _UpperCAmelCase   	=					unet_state_dict[
                  F'''output_blocks.{i}.{index}.conv.bias'''
              ]
  
              # Clear attentions as they have been attributed above.
              if len(a_				) == 2:
                 _UpperCAmelCase   	=					[]
  
           if len(a_				):
              _UpperCAmelCase   	=					renew_attention_paths(a_				)
              _UpperCAmelCase   	=					{
                  "old": F'''output_blocks.{i}.1''',
                  "new": F'''up_blocks.{block_id}.attentions.{layer_in_block_id}''',
              }
              assign_to_checkpoint(
                  a_    ,  a_    ,  a_    ,  additional_replacements=[meta_path]    ,  config=a_				)
  
           if len(a_				):
              _UpperCAmelCase   	=					renew_attention_paths(a_				)
              _UpperCAmelCase   	=					{
                  "old": F'''output_blocks.{i}.2''',
                  "new": F'''up_blocks.{block_id}.temp_attentions.{layer_in_block_id}''',
              }
              assign_to_checkpoint(
                  a_    ,  a_    ,  a_    ,  additional_replacements=[meta_path]    ,  config=a_				)
        else:
           _UpperCAmelCase   	=					renew_resnet_paths(a_    ,  n_shave_prefix_segments=1				)
           for path in resnet_0_paths:
              _UpperCAmelCase   	=					".".join(["output_blocks", str(a_				), path["old"]]				)
              _UpperCAmelCase   	=					".".join(["up_blocks", str(a_				), "resnets", str(a_				), path["new"]]				)
              _UpperCAmelCase   	=					unet_state_dict[old_path]
  
           _UpperCAmelCase   	=					[l for l in output_block_layers if "temopral_conv" in l]
           for path in temopral_conv_paths:
              _UpperCAmelCase   	=					path.split("temopral_conv."				)[-1]
              _UpperCAmelCase   	=					".".join(["output_blocks", str(a_				), str(a_				), "temopral_conv", pruned_path]				)
              _UpperCAmelCase   	=					".".join(["up_blocks", str(a_				), "temp_convs", str(a_				), pruned_path]				)
              _UpperCAmelCase   	=					unet_state_dict[old_path]
  
     return new_checkpoint
  
  
  if __name__ == "__main__":
    SCREAMING_SNAKE_CASE_     				=   argparse.ArgumentParser()
  
    parser.add_argument(
        '--checkpoint_path', default=None, type=str, required=True, help='Path to the checkpoint to convert.'
    )
    parser.add_argument('--dump_path', default=None, type=str, required=True, help='Path to the output model.')
    SCREAMING_SNAKE_CASE_     				=   parser.parse_args()
  
    SCREAMING_SNAKE_CASE_     				=   torch.load(args.checkpoint_path, map_location='cpu')
    SCREAMING_SNAKE_CASE_     				=   UNetaDConditionModel()
  
    SCREAMING_SNAKE_CASE_     				=   convert_ldm_unet_checkpoint(unet_checkpoint, unet.config)
  
    SCREAMING_SNAKE_CASE_     				=   set(unet.state_dict().keys()) - set(converted_ckpt.keys())
    SCREAMING_SNAKE_CASE_     				=   set(converted_ckpt.keys()) - set(unet.state_dict().keys())
  
    assert len(diff_a) == len(diff_a) == 0, "Converted weights don't match"
  
    # load state_dict
    unet.load_state_dict(converted_ckpt)
  
    unet.save_pretrained(args.dump_path)
  
    # -- finish converting the unet --
  
  
style_context: 
  """simple docstring"""
  
  
  
  
  
  
  
  from __future__ import annotations
  
  import unittest
  
  from transformers import AlbertConfig, is_tf_available
  from transformers.models.auto import get_values
  from transformers.testing_utils import require_tf, slow
  
  from ...test_configuration_common import ConfigTester
  from ...test_modeling_tf_common import TFModelTesterMixin, ids_tensor, random_attention_mask
  from ...test_pipeline_mixin import PipelineTesterMixin
  
  
  if is_tf_available():
    import tensorflow as tf
  
    from transformers import TF_MODEL_FOR_PRETRAINING_MAPPING
    from transformers.models.albert.modeling_tf_albert import (
        TF_ALBERT_PRETRAINED_MODEL_ARCHIVE_LIST,
        TFAlbertForMaskedLM,
        TFAlbertForMultipleChoice,
        TFAlbertForPreTraining,
        TFAlbertForQuestionAnswering,
        TFAlbertForSequenceClassification,
        TFAlbertForTokenClassification,
        TFAlbertModel,
    )
  
  
  
  
  
  
  class     __a			:
  
  
  
  
  
  
  
         def __init__(					self							,lowerCamelCase_							,lowerCamelCase_=13							,lowerCamelCase_=7							,lowerCamelCase_=True							,lowerCamelCase_=True							,lowerCamelCase_=True							,lowerCamelCase_=True							,lowerCamelCase_=99							,lowerCamelCase_=16							,lowerCamelCase_=32							,lowerCamelCase_=2							,lowerCamelCase_=4							,lowerCamelCase_=37							,lowerCamelCase_="gelu"							,lowerCamelCase_=0.1							,lowerCamelCase_=0.1							,lowerCamelCase_=512							,lowerCamelCase_=16							,lowerCamelCase_=2							,lowerCamelCase_=0.0_2							,lowerCamelCase_=3							,lowerCamelCase_=4							,lowerCamelCase_=None							,):
  
  
  
            '''simple docstring'''
  
            _UpperCAmelCase   	=					parent
            _UpperCAmelCase   	=					13
            _UpperCAmelCase   	=					7
            _UpperCAmelCase   	=					True
            _UpperCAmelCase   	=					True
            _UpperCAmelCase   	=					True
            _UpperCAmelCase   	=					True
            _UpperCAmelCase   	=					99
            _UpperCAmelCase   	=					16
            _UpperCAmelCase   	=					32
            _UpperCAmelCase   	=					2
            _UpperCAmelCase   	=					4
            _UpperCAmelCase   	=					37
            _UpperCAmelCase   	=					"gelu"
            _UpperCAmelCase   	=					0.1
            _UpperCAmelCase   	=					0.1
            _UpperCAmelCase   	=					512
            _UpperCAmelCase   	=					16
            _UpperCAmelCase   	=					2
            _UpperCAmelCase   	=					0.0_2
            _UpperCAmelCase   	=					3
            _UpperCAmelCase   	=					4
            _UpperCAmelCase   	=					None
  
  
  
  
  
  
  
         def    _snake_case			(					self		):
  
  
  
            '''simple docstring'''
  
            _UpperCAmelCase   	=					ids_tensor([self.batch_size, self.seq_length]							,self.vocab_size		)
  
            _UpperCAmelCase   	=					None
            if self.use_input_mask:
               _UpperCAmelCase   	=					random_attention_mask([self.batch_size, self.seq_length]		)
  
            _UpperCAmelCase   	=					None
            if self.use_token_type_ids:
               _UpperCAmelCase   	=					ids_tensor([self.batch_size, self.seq_length]							,self.type_vocab_size		)
  
            _UpperCAmelCase   	=					None
            _UpperCAmelCase   	=					None
            _UpperCAmelCase   	=					None
            if self.use_labels:
               _UpperCAmelCase   	=					ids_tensor([self.batch_size]							,self.type_sequence_label_size		)
               _UpperCAmelCase   	=					ids_tensor([self.batch_size, self.seq_length]							,self.num_labels		)
               _UpperCAmelCase   	=					ids_tensor([self.batch_size]							,self.num_choices		)
  
            _UpperCAmelCase   	=					AlbertConfig(
                vocab_size=self.vocab_size							,hidden_size=self.hidden_size							,embedding_size=self.embedding_size							,num_hidden_layers=self.num_hidden_layers							,num_attention_heads=self.num_attention_heads							,intermediate_size=self.intermediate_size							,hidden_act=self.hidden_act							,hidden_dropout_prob=self.hidden_dropout_prob							,attention_probs_dropout_prob=self.attention_probs_dropout_prob							,max_position_embeddings=self.max_position_embeddings							,type_vocab_size=self.type_vocab_size							,initializer_range=self.initializer_range							,)
  
            return config, input_ids, token_type_ids, input_mask, sequence_labels, token_labels, choice_labels
  
  
  
  
  
  
  
         def    _snake_case			(					self							,lowerCamelCase_							,lowerCamelCase_							,lowerCamelCase_							,lowerCamelCase_							,lowerCamelCase_							,lowerCamelCase_							,lowerCamelCase_		):
  
  
  
            '''simple docstring'''
  
            _UpperCAmelCase   	=					TFAlbertModel(config=lowerCamelCase_		)
            # inputs = {'input_ids': input_ids,
            #           'attention_mask': input_mask,
            #           'token_type_ids': token_type_ids}
            # sequence_output, pooled_output = model(**inputs)
            _UpperCAmelCase   	=					{"input_ids": input_ids, "attention_mask": input_mask, "token_type_ids": token_type_ids}
            _UpperCAmelCase   	=					model(lowerCamelCase_		)
  
            _UpperCAmelCase   	=					[input_ids, input_mask]
            _UpperCAmelCase   	=					model(lowerCamelCase_		)
  
            _UpperCAmelCase   	=					model(lowerCamelCase_		)
  
            self.parent.assertEqual(result.last_hidden_state.shape							,(self.batch_size, self.seq_length, self.hidden_size)		)
            self.parent.assertEqual(result.pooler_output.shape							,(self.batch_size, self.hidden_size)		)
  
  
  
  
  
  
  
         def    _snake_case			(					self							,lowerCamelCase_							,lowerCamelCase_							,lowerCamelCase_							,lowerCamelCase_							,lowerCamelCase_							,lowerCamelCase_							,lowerCamelCase_		):
  
  
  
            '''simple docstring'''
  
            _UpperCAmelCase   	=					self.num_labels
            _UpperCAmelCase   	=					TFAlbertForPreTraining(config=lowerCamelCase_		)
            _UpperCAmelCase   	=					{"input_ids": input_ids, "attention_mask": input_mask, "token_type_ids": token_type_ids}
            _UpperCAmelCase   	=					model(lowerCamelCase_		)
            self.parent.assertEqual(result.prediction_logits.shape							,(self.batch_size, self.seq_length, self.vocab_size)		)
            self.parent.assertEqual(result.sop_logits.shape							,(self.batch_size, self.num_labels)		)
  
  
  
  
  
  
  
         def    _snake_case			(					self							,lowerCamelCase_							,lowerCamelCase_							,lowerCamelCase_							,lowerCamelCase_							,lowerCamelCase_							,lowerCamelCase_							,lowerCamelCase_		):
  
  
  
            '''simple docstring'''
  
            _UpperCAmelCase   	=					TFAlbertForMaskedLM(config=lowerCamelCase_		)
            _UpperCAmelCase   	=					{"input_ids": input_ids, "attention_mask": input_mask, "token_type_ids": token_type_ids}
            _UpperCAmelCase   	=					model(lowerCamelCase_		)
            self.parent.assertEqual(result.logits.shape							,(self.batch_size, self.seq_length, self.vocab_size)		)
  
  
  
  
  
  
  
         def    _snake_case			(					self							,lowerCamelCase_							,lowerCamelCase_							,lowerCamelCase_							,lowerCamelCase_							,lowerCamelCase_							,lowerCamelCase_							,lowerCamelCase_		):
  
  
  
            '''simple docstring'''
  
            _UpperCAmelCase   	=					self.num_labels
            _UpperCAmelCase   	=					TFAlbertForSequenceClassification(config=lowerCamelCase_		)
            _UpperCAmelCase   	=					{"input_ids": input_ids, "attention_mask": input_mask, "token_type_ids": token_type_ids}
            _UpperCAmelCase   	=					model(lowerCamelCase_		)
            self.parent.assertEqual(result.logits.shape							,(self.batch_size, self.num_labels)		)
  
  
  
  
  
  
  
         def    _snake_case			(					self							,lowerCamelCase_							,lowerCamelCase_							,lowerCamelCase_							,lowerCamelCase_							,lowerCamelCase_							,lowerCamelCase_							,lowerCamelCase_		):
  
  
  
            '''simple docstring'''
  
            _UpperCAmelCase   	=					TFAlbertForQuestionAnswering(config=lowerCamelCase_		)
            _UpperCAmelCase   	=					{"input_ids": input_ids, "attention_mask": input_mask, "token_type_ids": token_type_ids}
            _UpperCAmelCase   	=					model(lowerCamelCase_		)
            self.parent.assertEqual(result.start_logits.shape							,(self.batch_size, self.seq_length)		)
            self.parent.assertEqual(result.end_logits.shape							,(self.batch_size, self.seq_length)		)
  
  
  
  
  
  
  
         def    _snake_case			(					self							,lowerCamelCase_							,lowerCamelCase_							,lowerCamelCase_							,lowerCamelCase_							,lowerCamelCase_							,lowerCamelCase_							,lowerCamelCase_		):
  
  
  
            '''simple docstring'''
  
            _UpperCAmelCase   	=					self.num_choices
            _UpperCAmelCase   	=					TFAlbertForMultipleChoice(config=lowerCamelCase_		)
            _UpperCAmelCase   	=					tf.tile(tf.expand_dims(lowerCamelCase_							,1		)							,(1, self.num_choices, 1)		)
            _UpperCAmelCase   	=					tf.tile(tf.expand_dims(lowerCamelCase_							,1		)							,(1, self.num_choices, 1)		)
            _UpperCAmelCase   	=					tf.tile(tf.expand_dims(lowerCamelCase_							,1		)							,(1, self.num_choices, 1)		)
            _UpperCAmelCase   	=					{
                "input_ids": multiple_choice_inputs_ids,
                "attention_mask": multiple_choice_input_mask,
                "token_type_ids": multiple_choice_token_type_ids,
            }
            _UpperCAmelCase   	=					model(lowerCamelCase_		)
            self.parent.assertListEqual(list(result["logits"].shape		)							,[self.batch_size, self.num_choices]		)
  
  
  
  
  
  
  
         def    _snake_case			(					self							,lowerCamelCase_							,lowerCamelCase_							,lowerCamelCase_							,lowerCamelCase_							,lowerCamelCase_							,lowerCamelCase_							,lowerCamelCase_		):
  
  
  
            '''simple docstring'''
  
            _UpperCAmelCase   	=					self.num_labels
            _UpperCAmelCase   	=					TFAlbertForTokenClassification(config=lowerCamelCase_		)
            _UpperCAmelCase   	=					{
                "input_ids": input_ids,
                "attention_mask": input_mask,
                "token_type_ids": token_type_ids,
            }
            _UpperCAmelCase   	=					model(lowerCamelCase_		)
            self.parent.assertListEqual(list(result["logits"].shape		)							,[self.batch_size, self.seq_length, self.num_labels]		)
  
  
  
  
  
  
  
         def    _snake_case			(					self		):
  
  
  
            '''simple docstring'''
  
            _UpperCAmelCase   	=					self.prepare_config_and_inputs()
            (
                (
                _UpperCAmelCase
            )	,							(
                _UpperCAmelCase
            )	,							(
                _UpperCAmelCase
            )	,							(
                _UpperCAmelCase
            )	,							(
                _UpperCAmelCase
            )	,							(
                _UpperCAmelCase
            )	,							(
                _UpperCAmelCase
            )	,							
            )   	=					config_and_inputs
            _UpperCAmelCase   	=					{"input_ids": input_ids, "token_type_ids": token_type_ids, "attention_mask": input_mask}
            return config, inputs_dict
  
  
  
  
  
  
  @require_tf
  class     __a			(					lowerCAmelCase_				,	lowerCAmelCase_				,	unittest.TestCase			):
         _lowerCamelCase							:							Any       =			(
             (
                 TFAlbertModel,
                 TFAlbertForPreTraining,
                 TFAlbertForMaskedLM,
                 TFAlbertForSequenceClassification,
                 TFAlbertForQuestionAnswering,
                 TFAlbertForTokenClassification,
                 TFAlbertForMultipleChoice,
             )
             if is_tf_available()
             else ()
         )
         _lowerCamelCase							:							int       =			(
             {
                 'feature-extraction': TFAlbertModel,
                 'fill-mask': TFAlbertForMaskedLM,
                 'question-answering': TFAlbertForQuestionAnswering,
                 'text-classification': TFAlbertForSequenceClassification,
                 'token-classification': TFAlbertForTokenClassification,
                 'zero-shot': TFAlbertForSequenceClassification,
             }
             if is_tf_available()
             else {}
         )
         _lowerCamelCase							:							List[str]       =			False
         _lowerCamelCase							:							Optional[Any]       =			False
  
  
  
  
  
  
  
         def    _snake_case			(					self							,lowerCamelCase_							,lowerCamelCase_							,lowerCamelCase_=False		):
  
  
  
            '''simple docstring'''
  
            _UpperCAmelCase   	=					super()._prepare_for_class(lowerCamelCase_							,lowerCamelCase_							,return_labels=lowerCamelCase_		)
  
            if return_labels:
               if model_class in get_values(lowerCamelCase_		):
                  _UpperCAmelCase   	=					tf.zeros(self.model_tester.batch_size							,dtype=tf.intaa		)
  
            return inputs_dict
  
  
  
  
  
  
  
         def    _snake_case			(					self		):
  
  
  
            '''simple docstring'''
  
            _UpperCAmelCase   	=					TFAlbertModelTester(self		)
            _UpperCAmelCase   	=					ConfigTester(self							,config_class=lowerCamelCase_							,hidden_size=37		)
  
  
  
  
  
  
  
         def    _snake_case			(					self		):
  
  
  
            '''simple docstring'''
  
            self.config_tester.run_common_tests()
  
  
  
  
  
  
  
         def    _snake_case			(					self		):
  
  
  
            '''simple docstring'''
  
            _UpperCAmelCase   	=					self.model_tester.prepare_config_and_inputs()
            self.model_tester.create_and_check_albert_model(*lowerCamelCase_		)
  
  
  
  
  
  
  
         def    _snake_case			(					self		):
  
  
  
            '''simple docstring'''
  
            _UpperCAmelCase   	=					self.model_tester.prepare_config_and_inputs()
            self.model_tester.create_and_check_albert_for_pretraining(*lowerCamelCase_		)
  
  
  
  
  
  
  
         def    _snake_case			(					self		):
  
  
  
            '''simple docstring'''
  
            _UpperCAmelCase   	=					self.model_tester.prepare_config_and_inputs()
            self.model_tester.create_and_check_albert_for_masked_lm(*lowerCamelCase_		)
  
  
  
  
  
  
  
         def    _snake_case			(					self		):
  
  
  
            '''simple docstring'''
  
            _UpperCAmelCase   	=					self.model_tester.prepare_config_and_inputs()
            self.model_tester.create_and_check_albert_for_multiple_choice(*lowerCamelCase_		)
  
  
  
  
  
  
  
         def    _snake_case			(					self		):
  
  
  
            '''simple docstring'''
  
            _UpperCAmelCase   	=					self.model_tester.prepare_config_and_inputs()
            self.model_tester.create_and_check_albert_for_sequence_classification(*lowerCamelCase_		)
  
  
  
  
  
  
  
         def    _snake_case			(					self		):
  
  
  
            '''simple docstring'''
  
            _UpperCAmelCase   	=					self.model_tester.prepare_config_and_inputs()
            self.model_tester.create_and_check_albert_for_question_answering(*lowerCamelCase_		)
  
  
  
  
  
  
  
         @slow
         def    _snake_case			(					self		):
  
  
  
            '''simple docstring'''
  
            for model_name in TF_ALBERT_PRETRAINED_MODEL_ARCHIVE_LIST[:1]:
               _UpperCAmelCase   	=					TFAlbertModel.from_pretrained(lowerCamelCase_		)
               self.assertIsNotNone(lowerCamelCase_		)
  
  
  
  
  
  
  @require_tf
  class     __a			(					unittest.TestCase			):
  
  
  
  
  
  
  
         @slow
         def    _snake_case			(					self		):
  
  
  
            '''simple docstring'''
  
            _UpperCAmelCase   	=					TFAlbertForPreTraining.from_pretrained("albert-base-v2"		)
            _UpperCAmelCase   	=					tf.constant([[0, 1, 2, 3, 4, 5]]		)
            _UpperCAmelCase   	=					model(lowerCamelCase_		)[0]
  
            _UpperCAmelCase   	=					[1, 6, 3_0000]
            self.assertEqual(output.shape							,lowerCamelCase_		)
  
            _UpperCAmelCase   	=					tf.constant(
                [
                    [
                        [4.5_9_5_6_6_8, 0.7_4_4_6_2_7_5_4, -1.8_1_8_1_4_7],
                        [4.5_9_5_4_3_4_7, 0.7_4_5_4_1_8_4, -1.8_1_8_8_2_5_8],
                        [4.5_9_5_4_9_0_5, 0.7_4_4_8_2_3_5, -1.8_1_8_2_3_1_6],
                    ]
                ]		)
            tf.debugging.assert_near(output[:, :3, :3]							,lowerCamelCase_							,atol=1e-4		)
  
  