import inspect
import os

import torch
from transformers import AutoModel
from transformers.testing_utils import mockenv_context
from transformers.trainer_utils import set_seed

import accelerate
from accelerate.accelerator import Accelerator
from accelerate.state import AcceleratorState
from accelerate.test_utils.testing import (
    AccelerateTestCase,
    TempDirTestCase,
    execute_subprocess_async,
    require_cuda,
    require_fsdp,
    require_multi_gpu,
    slow,
)
from accelerate.utils.constants import (
    FSDP_AUTO_WRAP_POLICY,
    FSDP_BACKWARD_PREFETCH,
    FSDP_SHARDING_STRATEGY,
    FSDP_STATE_DICT_TYPE,
)
from accelerate.utils.dataclasses import FullyShardedDataParallelPlugin
from accelerate.utils.other import patch_environment


set_seed(4_2)

UpperCAmelCase_           =				'''bert-base-cased'''
UpperCAmelCase_           =				'''fp16'''
UpperCAmelCase_           =				'''bf16'''
UpperCAmelCase_           =				[FPaa, BFaa]




@require_fsdp
@require_cuda
class 			__magic_name__	(		__a     ):





   """simple docstring"""




   def 	lowerCAmelCase    (						self       :	Dict				):



    """simple docstring"""



    super().setUp()

    _UpperCamelCase:  List[Any]    						=			dict(
        ACCELERATE_USE_FSDP='''true'''		,	MASTER_ADDR='''localhost'''		,	MASTER_PORT='''10999'''		,	RANK='''0'''		,	LOCAL_RANK='''0'''		,	WORLD_SIZE='''1'''		,	)



   def 	lowerCAmelCase    (						self       :	List[Any]				):



    """simple docstring"""



    from torch.distributed.fsdp.fully_sharded_data_parallel import ShardingStrategy

    for i, strategy in enumerate(_lowercase				):
     _UpperCamelCase:  Optional[int]    						=			self.dist_env.copy()
     _UpperCamelCase:  List[str]    						=			f"""{i + 1}"""
     _UpperCamelCase:  Dict    						=			strategy
     with mockenv_context(**_lowercase				):
      _UpperCamelCase:  List[str]    						=			FullyShardedDataParallelPlugin()
      self.assertEqual(fsdp_plugin.sharding_strategy		,	ShardingStrategy(i + 1				)				)



   def 	lowerCAmelCase    (						self       :	Union[str, Any]				):



    """simple docstring"""



    from torch.distributed.fsdp.fully_sharded_data_parallel import BackwardPrefetch

    for i, prefetch_policy in enumerate(_lowercase				):
     _UpperCamelCase:  Any    						=			self.dist_env.copy()
     _UpperCamelCase:  Any    						=			prefetch_policy
     with mockenv_context(**_lowercase				):
      _UpperCamelCase:  Union[str, Any]    						=			FullyShardedDataParallelPlugin()
      if prefetch_policy == "NO_PREFETCH":
       self.assertIsNone(fsdp_plugin.backward_prefetch				)
      else:
       self.assertEqual(fsdp_plugin.backward_prefetch		,	BackwardPrefetch(i + 1				)				)



   def 	lowerCAmelCase    (						self       :	Dict				):



    """simple docstring"""



    from torch.distributed.fsdp.fully_sharded_data_parallel import StateDictType

    for i, state_dict_type in enumerate(_lowercase				):
     _UpperCamelCase:  Dict    						=			self.dist_env.copy()
     _UpperCamelCase:  str    						=			state_dict_type
     with mockenv_context(**_lowercase				):
      _UpperCamelCase:  Any    						=			FullyShardedDataParallelPlugin()
      self.assertEqual(fsdp_plugin.state_dict_type		,	StateDictType(i + 1				)				)
      if state_dict_type == "FULL_STATE_DICT":
       self.assertTrue(fsdp_plugin.state_dict_config.offload_to_cpu				)
       self.assertTrue(fsdp_plugin.state_dict_config.ranka_only				)



   def 	lowerCAmelCase    (						self       :	Optional[int]				):



    """simple docstring"""



    _UpperCamelCase:  List[Any]    						=			AutoModel.from_pretrained(_lowercase				)
    for policy in FSDP_AUTO_WRAP_POLICY:
     _UpperCamelCase:  str    						=			self.dist_env.copy()
     _UpperCamelCase:  str    						=			policy
     if policy == "TRANSFORMER_BASED_WRAP":
      _UpperCamelCase:  Optional[Any]    						=			'''BertLayer'''
     elif policy == "SIZE_BASED_WRAP":
      _UpperCamelCase:  Tuple    						=			'''2000'''
     with mockenv_context(**_lowercase				):
      _UpperCamelCase:  Union[str, Any]    						=			FullyShardedDataParallelPlugin()
      fsdp_plugin.set_auto_wrap_policy(_lowercase				)
      if policy == "NO_WRAP":
       self.assertIsNone(fsdp_plugin.auto_wrap_policy				)
      else:
       self.assertIsNotNone(fsdp_plugin.auto_wrap_policy				)

    _UpperCamelCase:  Optional[int]    						=			self.dist_env.copy()
    _UpperCamelCase:  int    						=			'''TRANSFORMER_BASED_WRAP'''
    _UpperCamelCase:  List[Any]    						=			'''T5Layer'''
    with mockenv_context(**_lowercase				):
     _UpperCamelCase:  str    						=			FullyShardedDataParallelPlugin()
     with self.assertRaises(_lowercase				) as cm:
      fsdp_plugin.set_auto_wrap_policy(_lowercase				)
     self.assertTrue('''Could not find the transformer layer class to wrap in the model.''' in str(cm.exception				)				)

    _UpperCamelCase:  str    						=			self.dist_env.copy()
    _UpperCamelCase:  Tuple    						=			'''SIZE_BASED_WRAP'''
    _UpperCamelCase:  Optional[int]    						=			'''0'''
    with mockenv_context(**_lowercase				):
     _UpperCamelCase:  List[str]    						=			FullyShardedDataParallelPlugin()
     fsdp_plugin.set_auto_wrap_policy(_lowercase				)
     self.assertIsNone(fsdp_plugin.auto_wrap_policy				)



   def 	lowerCAmelCase    (						self       :	List[Any]				):



    """simple docstring"""



    from torch.distributed.fsdp.fully_sharded_data_parallel import MixedPrecision
    from torch.distributed.fsdp.sharded_grad_scaler import ShardedGradScaler

    for mp_dtype in dtypes:
     _UpperCamelCase:  int    						=			self.dist_env.copy()
     _UpperCamelCase:  Dict    						=			mp_dtype
     with mockenv_context(**_lowercase				):
      _UpperCamelCase:  Union[str, Any]    						=			Accelerator()
      if mp_dtype == "fp16":
       _UpperCamelCase:  Optional[Any]    						=			torch.floataa
      elif mp_dtype == "bf16":
       _UpperCamelCase:  int    						=			torch.bfloataa
      _UpperCamelCase:  Any    						=			MixedPrecision(param_dtype=_lowercase		,	reduce_dtype=_lowercase		,	buffer_dtype=_lowercase				)
      self.assertEqual(accelerator.state.fsdp_plugin.mixed_precision_policy		,	_lowercase				)
      if mp_dtype == FPaa:
       self.assertTrue(isinstance(accelerator.scaler		,	_lowercase				)				)
      elif mp_dtype == BFaa:
       self.assertIsNone(accelerator.scaler				)
      AcceleratorState._reset_state(_lowercase				)



   def 	lowerCAmelCase    (						self       :	Union[str, Any]				):



    """simple docstring"""



    from torch.distributed.fsdp.fully_sharded_data_parallel import CPUOffload

    for flag in [True, False]:
     _UpperCamelCase:  Union[str, Any]    						=			self.dist_env.copy()
     _UpperCamelCase:  Dict    						=			str(_lowercase				).lower()
     with mockenv_context(**_lowercase				):
      _UpperCamelCase:  List[Any]    						=			FullyShardedDataParallelPlugin()
      self.assertEqual(fsdp_plugin.cpu_offload		,	CPUOffload(offload_params=_lowercase				)				)




@require_fsdp
@require_multi_gpu
@slow
class 			__magic_name__	(		__a     ):





   """simple docstring"""




   def 	lowerCAmelCase    (						self       :	Any				):



    """simple docstring"""



    super().setUp()
    _UpperCamelCase:  List[Any]    						=			0.82
    _UpperCamelCase:  Optional[Any]    						=			[
        '''fsdp_shard_grad_op_transformer_based_wrap''',
        '''fsdp_full_shard_transformer_based_wrap''',
    ]
    _UpperCamelCase:  Dict    						=			{
        '''multi_gpu_fp16''': 3_200,
        '''fsdp_shard_grad_op_transformer_based_wrap_fp16''': 2_000,
        '''fsdp_full_shard_transformer_based_wrap_fp16''': 1_900,
        # Disabling below test as it overwhelms the RAM memory usage
        # on CI self-hosted runner leading to tests getting killed.
        # "fsdp_full_shard_cpu_offload_transformer_based_wrap_fp32": 1500,  # fp16 was leading to indefinite hang
    }
    _UpperCamelCase:  int    						=			160
    _UpperCamelCase:  Any    						=			160

    _UpperCamelCase:  Optional[Any]    						=			inspect.getfile(accelerate.test_utils				)
    _UpperCamelCase:  List[Any]    						=			os.path.sep.join(mod_file.split(os.path.sep				)[:-1] + ['''scripts''', '''external_deps''']				)



   def 	lowerCAmelCase    (						self       :	Optional[Any]				):



    """simple docstring"""



    _UpperCamelCase:  Any    						=			os.path.join(self.test_scripts_folder		,	'''test_performance.py'''				)
    _UpperCamelCase:  Optional[int]    						=			['''accelerate''', '''launch''', '''--num_processes=2''', '''--num_machines=1''', '''--machine_rank=0''', '''--use_fsdp''']
    for config in self.performance_configs:
     _UpperCamelCase:  List[str]    						=			cmd.copy()
     for i, strategy in enumerate(_lowercase				):
      if strategy.lower() in config:
       cmd_config.append(f"""--fsdp_sharding_strategy={i+1}"""				)
       break

     if "fp32" in config:
      cmd_config.append('''--mixed_precision=no'''				)
     else:
      cmd_config.append('''--mixed_precision=fp16'''				)

     if "cpu_offload" in config:
      cmd_config.append('''--fsdp_offload_params=True'''				)

     for policy in FSDP_AUTO_WRAP_POLICY:
      if policy.lower() in config:
       cmd_config.append(f"""--fsdp_auto_wrap_policy={policy}"""				)
       break

     if policy == "TRANSFORMER_BASED_WRAP":
      cmd_config.append('''--fsdp_transformer_layer_cls_to_wrap=BertLayer'''				)
     elif policy == "SIZE_BASED_WRAP":
      cmd_config.append('''--fsdp_min_num_params=2000'''				)

     cmd_config.extend(
         [
             self.test_file_path,
             f"""--output_dir={self.tmpdir}""",
             f"""--performance_lower_bound={self.performance_lower_bound}""",
         ]				)
     with patch_environment(omp_num_threads=1				):
      execute_subprocess_async(_lowercase		,	env=os.environ.copy()				)



   def 	lowerCAmelCase    (						self       :	Union[str, Any]				):



    """simple docstring"""



    _UpperCamelCase:  Dict    						=			os.path.join(self.test_scripts_folder		,	'''test_checkpointing.py'''				)
    _UpperCamelCase:  int    						=			[
        '''accelerate''',
        '''launch''',
        '''--num_processes=2''',
        '''--num_machines=1''',
        '''--machine_rank=0''',
        '''--use_fsdp''',
        '''--mixed_precision=fp16''',
        '''--fsdp_transformer_layer_cls_to_wrap=BertLayer''',
    ]

    for i, strategy in enumerate(_lowercase				):
     _UpperCamelCase:  Any    						=			cmd.copy()
     cmd_config.append(f"""--fsdp_sharding_strategy={i+1}"""				)
     if strategy != "FULL_SHARD":
      continue
     _UpperCamelCase:  Union[str, Any]    						=			len(_lowercase				)
     for state_dict_type in FSDP_STATE_DICT_TYPE:
      _UpperCamelCase:  Optional[Any]    						=			cmd_config[:state_dict_config_index]
      cmd_config.append(f"""--fsdp_state_dict_type={state_dict_type}"""				)
      cmd_config.extend(
          [
              self.test_file_path,
              f"""--output_dir={self.tmpdir}""",
              '''--partial_train_epoch=1''',
          ]				)
      with patch_environment(omp_num_threads=1				):
       execute_subprocess_async(_lowercase		,	env=os.environ.copy()				)

      _UpperCamelCase:  List[Any]    						=			cmd_config[:-1]
      _UpperCamelCase:  Optional[Any]    						=			os.path.join(self.tmpdir		,	'''epoch_0'''				)
      cmd_config.extend(
          [
              f"""--resume_from_checkpoint={resume_from_checkpoint}""",
          ]				)
      with patch_environment(omp_num_threads=1				):
       execute_subprocess_async(_lowercase		,	env=os.environ.copy()				)



   def 	lowerCAmelCase    (						self       :	Optional[Any]				):



    """simple docstring"""



    _UpperCamelCase:  Dict    						=			os.path.join(self.test_scripts_folder		,	'''test_peak_memory_usage.py'''				)
    _UpperCamelCase:  List[str]    						=			[
        '''accelerate''',
        '''launch''',
        '''--num_processes=2''',
        '''--num_machines=1''',
        '''--machine_rank=0''',
    ]
    for spec, peak_mem_upper_bound in self.peak_memory_usage_upper_bound.items():
     _UpperCamelCase:  int    						=			cmd.copy()
     if "fp16" in spec:
      cmd_config.extend(['''--mixed_precision=fp16''']				)
     else:
      cmd_config.extend(['''--mixed_precision=no''']				)

     if "multi_gpu" in spec:
      continue
     else:
      cmd_config.extend(['''--use_fsdp''']				)
      for i, strategy in enumerate(_lowercase				):
       if strategy.lower() in spec:
        cmd_config.append(f"""--fsdp_sharding_strategy={i+1}"""				)
        break

      if "cpu_offload" in spec:
       cmd_config.append('''--fsdp_offload_params=True'''				)

      for policy in FSDP_AUTO_WRAP_POLICY:
       if policy.lower() in spec:
        cmd_config.append(f"""--fsdp_auto_wrap_policy={policy}"""				)
        break

      if policy == "TRANSFORMER_BASED_WRAP":
       cmd_config.append('''--fsdp_transformer_layer_cls_to_wrap=BertLayer'''				)
      elif policy == "SIZE_BASED_WRAP":
       cmd_config.append('''--fsdp_min_num_params=2000'''				)

     cmd_config.extend(
         [
             self.test_file_path,
             f"""--output_dir={self.tmpdir}""",
             f"""--peak_memory_upper_bound={peak_mem_upper_bound}""",
             f"""--n_train={self.n_train}""",
             f"""--n_val={self.n_val}""",
         ]				)
     with patch_environment(omp_num_threads=1				):
      execute_subprocess_async(_lowercase		,	env=os.environ.copy()				)
from __future__ import annotations

import unittest

from transformers import LEDConfig, is_tf_available
from transformers.testing_utils import require_tf, slow

from ...test_configuration_common import ConfigTester
from ...test_modeling_tf_common import TFModelTesterMixin, ids_tensor
from ...test_pipeline_mixin import PipelineTesterMixin


if is_tf_available():
  import tensorflow as tf

  from transformers import TFLEDForConditionalGeneration, TFLEDModel




@require_tf
class 			__magic_name__	:





   """simple docstring"""




   lowerCAmelCase   :     Optional[int]	      = LEDConfig
   lowerCAmelCase   :     List[str]	      = {}
   lowerCAmelCase   :     str	      = '''gelu'''



   def __init__(						self       :	int		,	_lowercase       :	Tuple		,	_lowercase       :	Dict=13		,	_lowercase       :	Union[str, Any]=7		,	_lowercase       :	Optional[int]=True		,	_lowercase       :	Any=False		,	_lowercase       :	Optional[int]=99		,	_lowercase       :	Dict=32		,	_lowercase       :	Optional[Any]=2		,	_lowercase       :	Union[str, Any]=4		,	_lowercase       :	List[str]=37		,	_lowercase       :	Optional[Any]=0.1		,	_lowercase       :	List[Any]=0.1		,	_lowercase       :	Union[str, Any]=20		,	_lowercase       :	Tuple=2		,	_lowercase       :	str=1		,	_lowercase       :	Optional[int]=0		,	_lowercase       :	List[str]=4		,	):



    """simple docstring"""



    _UpperCamelCase:  Optional[Any]    						=			parent
    _UpperCamelCase:  Union[str, Any]    						=			batch_size
    _UpperCamelCase:  Optional[Any]    						=			seq_length
    _UpperCamelCase:  str    						=			is_training
    _UpperCamelCase:  Dict    						=			use_labels
    _UpperCamelCase:  List[Any]    						=			vocab_size
    _UpperCamelCase:  Dict    						=			hidden_size
    _UpperCamelCase:  int    						=			num_hidden_layers
    _UpperCamelCase:  Dict    						=			num_attention_heads
    _UpperCamelCase:  Dict    						=			intermediate_size

    _UpperCamelCase:  Union[str, Any]    						=			hidden_dropout_prob
    _UpperCamelCase:  Any    						=			attention_probs_dropout_prob
    _UpperCamelCase:  Optional[Any]    						=			max_position_embeddings
    _UpperCamelCase:  Union[str, Any]    						=			eos_token_id
    _UpperCamelCase:  Union[str, Any]    						=			pad_token_id
    _UpperCamelCase:  Optional[Any]    						=			bos_token_id
    _UpperCamelCase:  Optional[Any]    						=			attention_window

    # `ModelTesterMixin.test_attention_outputs` is expecting attention tensors to be of size
    # [num_attention_heads, encoder_seq_length, encoder_key_length], but TFLongformerSelfAttention
    # returns attention of shape [num_attention_heads, encoder_seq_length, self.attention_window + 1]
    # because its local attention only attends to `self.attention_window` and one before and one after
    _UpperCamelCase:  Any    						=			self.attention_window + 2

    # because of padding `encoder_seq_length`, is different from `seq_length`. Relevant for
    # the `test_attention_outputs` and `test_hidden_states_output` tests
    _UpperCamelCase:  Tuple    						=			(
        self.seq_length + (self.attention_window - self.seq_length % self.attention_window) % self.attention_window
    )



   def 	lowerCAmelCase    (						self       :	int				):



    """simple docstring"""



    _UpperCamelCase:  Union[str, Any]    						=			ids_tensor([self.batch_size, self.seq_length - 1]		,	self.vocab_size				)
    _UpperCamelCase:  Union[str, Any]    						=			tf.expand_dims(tf.constant([self.eos_token_id] * self.batch_size				)		,	1				)
    _UpperCamelCase:  Union[str, Any]    						=			tf.concat([input_ids, eos_tensor]		,	axis=1				)

    _UpperCamelCase:  Union[str, Any]    						=			ids_tensor([self.batch_size, self.seq_length]		,	self.vocab_size				)

    _UpperCamelCase:  Any    						=			self.config_cls(
        vocab_size=self.vocab_size		,	d_model=self.hidden_size		,	encoder_layers=self.num_hidden_layers		,	decoder_layers=self.num_hidden_layers		,	encoder_attention_heads=self.num_attention_heads		,	decoder_attention_heads=self.num_attention_heads		,	encoder_ffn_dim=self.intermediate_size		,	decoder_ffn_dim=self.intermediate_size		,	dropout=self.hidden_dropout_prob		,	attention_dropout=self.attention_probs_dropout_prob		,	max_position_embeddings=self.max_position_embeddings		,	eos_token_ids=[2]		,	bos_token_id=self.bos_token_id		,	pad_token_id=self.pad_token_id		,	decoder_start_token_id=self.pad_token_id		,	attention_window=self.attention_window		,	**self.config_updates		,	)
    _UpperCamelCase:  Dict    						=			prepare_led_inputs_dict(_lowercase		,	_lowercase		,	_lowercase				)
    _UpperCamelCase:  Tuple    						=			tf.concat(
        [tf.zeros_like(_lowercase				)[:, :-1], tf.ones_like(_lowercase				)[:, -1:]]		,	axis=-1		,	)
    _UpperCamelCase:  List[Any]    						=			global_attention_mask
    return config, inputs_dict



   def 	lowerCAmelCase    (						self       :	Dict		,	_lowercase       :	List[str]		,	_lowercase       :	Union[str, Any]				):



    """simple docstring"""



    _UpperCamelCase:  Optional[Any]    						=			TFLEDModel(config=_lowercase				).get_decoder()
    _UpperCamelCase:  Optional[int]    						=			inputs_dict['''input_ids''']

    _UpperCamelCase:  Dict    						=			input_ids[:1, :]
    _UpperCamelCase:  Optional[Any]    						=			inputs_dict['''attention_mask'''][:1, :]
    _UpperCamelCase:  Optional[int]    						=			1

    # first forward pass
    _UpperCamelCase:  Optional[Any]    						=			model(_lowercase		,	attention_mask=_lowercase		,	use_cache=_lowercase				)

    _UpperCamelCase						,				_UpperCamelCase:  Union[str, Any]    						=			outputs.to_tuple()

    # create hypothetical next token and extent to next_input_ids
    _UpperCamelCase:  Any    						=			ids_tensor((self.batch_size, 3)		,	config.vocab_size				)
    _UpperCamelCase:  List[str]    						=			tf.cast(ids_tensor((self.batch_size, 3)		,	2				)		,	tf.inta				)

    # append to next input_ids and
    _UpperCamelCase:  Optional[int]    						=			tf.concat([input_ids, next_tokens]		,	axis=-1				)
    _UpperCamelCase:  Union[str, Any]    						=			tf.concat([attention_mask, next_attn_mask]		,	axis=-1				)

    _UpperCamelCase:  str    						=			model(_lowercase		,	attention_mask=_lowercase				)[0]
    _UpperCamelCase:  Optional[int]    						=			model(_lowercase		,	attention_mask=_lowercase		,	past_key_values=_lowercase				)[0]

    self.parent.assertEqual(next_tokens.shape[1]		,	output_from_past.shape[1]				)

    # select random slice
    _UpperCamelCase:  Dict    						=			int(ids_tensor((1,)		,	output_from_past.shape[-1]				)				)
    _UpperCamelCase:  Any    						=			output_from_no_past[:, -3:, random_slice_idx]
    _UpperCamelCase:  Optional[Any]    						=			output_from_past[:, :, random_slice_idx]

    # test that outputs are equal for slice
    tf.debugging.assert_near(_lowercase		,	_lowercase		,	rtol=1E-3				)




def    lowerCAmelCase_  (		lowercase:				Tuple	,       lowercase:				List[str]	,       lowercase:				List[str]	,       lowercase:				int=None	,       lowercase:				int=None	,       lowercase:				List[str]=None	,       lowercase:				List[str]=None	,       )     ->					List[str]:







 '''simple docstring'''
 if attention_mask is None:
  _UpperCamelCase:  Optional[Any]    						=			tf.cast(tf.math.not_equal(lowercase	,       config.pad_token_id		)	,       tf.inta		)
 if decoder_attention_mask is None:
  _UpperCamelCase:  int    						=			tf.concat(
      [
          tf.ones(decoder_input_ids[:, :1].shape	,       dtype=tf.inta		),
          tf.cast(tf.math.not_equal(decoder_input_ids[:, 1:]	,       config.pad_token_id		)	,       tf.inta		),
      ]	,       axis=-1	,       )
 if head_mask is None:
  _UpperCamelCase:  str    						=			tf.ones((config.encoder_layers, config.encoder_attention_heads)		)
 if decoder_head_mask is None:
  _UpperCamelCase:  Optional[Any]    						=			tf.ones((config.decoder_layers, config.decoder_attention_heads)		)
 return {
     "input_ids": input_ids,
     "attention_mask": attention_mask,
     "decoder_input_ids": decoder_input_ids,
     "decoder_attention_mask": decoder_attention_mask,
     "head_mask": head_mask,
     "decoder_head_mask": decoder_head_mask,
 }




@require_tf
class 			__magic_name__	(		__a   ,						__a   ,						unittest.TestCase     ):





   """simple docstring"""




   lowerCAmelCase   :     Optional[int]	      = (TFLEDForConditionalGeneration, TFLEDModel) if is_tf_available() else ()
   lowerCAmelCase   :     List[str]	      = (TFLEDForConditionalGeneration,) if is_tf_available() else ()
   lowerCAmelCase   :     Union[str, Any]	      = (
       {
           '''conversational''': TFLEDForConditionalGeneration,
           '''feature-extraction''': TFLEDModel,
           '''summarization''': TFLEDForConditionalGeneration,
           '''text2text-generation''': TFLEDForConditionalGeneration,
           '''translation''': TFLEDForConditionalGeneration,
       }
       if is_tf_available()
       else {}
   )
   lowerCAmelCase   :     str	      = True
   lowerCAmelCase   :     Any	      = False
   lowerCAmelCase   :     Union[str, Any]	      = False
   lowerCAmelCase   :     Any	      = False



   def 	lowerCAmelCase    (						self       :	Any				):



    """simple docstring"""



    _UpperCamelCase:  Dict    						=			TFLEDModelTester(self				)
    _UpperCamelCase:  Any    						=			ConfigTester(self		,	config_class=_lowercase				)



   def 	lowerCAmelCase    (						self       :	Optional[int]				):



    """simple docstring"""



    self.config_tester.run_common_tests()



   def 	lowerCAmelCase    (						self       :	Tuple				):



    """simple docstring"""



    _UpperCamelCase:  Union[str, Any]    						=			self.model_tester.prepare_config_and_inputs_for_common()
    self.model_tester.check_decoder_model_past_large_inputs(*_lowercase				)



   def 	lowerCAmelCase    (						self       :	Optional[Any]				):



    """simple docstring"""



    _UpperCamelCase						,				_UpperCamelCase:  str    						=			self.model_tester.prepare_config_and_inputs_for_common()
    _UpperCamelCase:  str    						=			tf.zeros_like(inputs_dict['''attention_mask''']				)
    _UpperCamelCase:  Any    						=			2
    _UpperCamelCase:  List[str]    						=			tf.where(
        tf.range(self.model_tester.seq_length				)[None, :] < num_global_attn_indices		,	1		,	inputs_dict['''global_attention_mask''']		,	)

    _UpperCamelCase:  int    						=			True
    _UpperCamelCase:  str    						=			self.model_tester.seq_length
    _UpperCamelCase:  Dict    						=			self.model_tester.encoder_seq_length

    def check_decoder_attentions_output(_lowercase       :	Union[str, Any]				):
     _UpperCamelCase:  Tuple    						=			outputs.decoder_attentions
     self.assertEqual(len(_lowercase				)		,	self.model_tester.num_hidden_layers				)
     self.assertListEqual(
         list(decoder_attentions[0].shape[-3:]				)		,	[self.model_tester.num_attention_heads, seq_length, seq_length]		,	)

    def check_encoder_attentions_output(_lowercase       :	int				):
     _UpperCamelCase:  str    						=			[t.numpy() for t in outputs.encoder_attentions]
     _UpperCamelCase:  Any    						=			[t.numpy() for t in outputs.encoder_global_attentions]
     self.assertEqual(len(_lowercase				)		,	self.model_tester.num_hidden_layers				)
     self.assertEqual(len(_lowercase				)		,	self.model_tester.num_hidden_layers				)
     self.assertListEqual(
         list(attentions[0].shape[-3:]				)		,	[self.model_tester.num_attention_heads, seq_length, seq_length]		,	)
     self.assertListEqual(
         list(global_attentions[0].shape[-3:]				)		,	[self.model_tester.num_attention_heads, encoder_seq_length, num_global_attn_indices]		,	)

    for model_class in self.all_model_classes:
     _UpperCamelCase:  Optional[int]    						=			True
     _UpperCamelCase:  List[Any]    						=			False
     _UpperCamelCase:  Optional[int]    						=			False
     _UpperCamelCase:  Union[str, Any]    						=			model_class(_lowercase				)
     _UpperCamelCase:  Optional[int]    						=			model(self._prepare_for_class(_lowercase		,	_lowercase				)				)
     _UpperCamelCase:  List[str]    						=			len(_lowercase				)
     self.assertEqual(config.output_hidden_states		,	_lowercase				)
     check_encoder_attentions_output(_lowercase				)

     if self.is_encoder_decoder:
      _UpperCamelCase:  Optional[int]    						=			model_class(_lowercase				)
      _UpperCamelCase:  Union[str, Any]    						=			model(self._prepare_for_class(_lowercase		,	_lowercase				)				)
      self.assertEqual(config.output_hidden_states		,	_lowercase				)
      check_decoder_attentions_output(_lowercase				)

     # Check that output attentions can also be changed via the config
     del inputs_dict["output_attentions"]
     _UpperCamelCase:  Union[str, Any]    						=			True
     _UpperCamelCase:  List[str]    						=			model_class(_lowercase				)
     _UpperCamelCase:  Dict    						=			model(self._prepare_for_class(_lowercase		,	_lowercase				)				)
     self.assertEqual(config.output_hidden_states		,	_lowercase				)
     check_encoder_attentions_output(_lowercase				)

     # Check attention is always last and order is fine
     _UpperCamelCase:  str    						=			True
     _UpperCamelCase:  Dict    						=			True
     _UpperCamelCase:  int    						=			model_class(_lowercase				)
     _UpperCamelCase:  Optional[int]    						=			model(self._prepare_for_class(_lowercase		,	_lowercase				)				)

     self.assertEqual(out_len + (2 if self.is_encoder_decoder else 1)		,	len(_lowercase				)				)
     self.assertEqual(model.config.output_hidden_states		,	_lowercase				)
     check_encoder_attentions_output(_lowercase				)



   @unittest.skip('''LED keeps using potentially symbolic tensors in conditionals and breaks tracing.'''				)
   def 	lowerCAmelCase    (						self       :	List[Any]				):



    """simple docstring"""



    pass



   def 	lowerCAmelCase    (						self       :	Optional[int]				):



    """simple docstring"""



    pass




def    lowerCAmelCase_  (		lowercase:				Tuple		)     ->					List[str]:







 '''simple docstring'''
 return tf.constant(lowercase	,       dtype=tf.intaa		)


UpperCAmelCase_           =				1E-4





@slow
@require_tf
class 			__magic_name__	(		unittest.TestCase     ):





   """simple docstring"""




   def 	lowerCAmelCase    (						self       :	str				):



    """simple docstring"""



    _UpperCamelCase:  List[str]    						=			TFLEDForConditionalGeneration.from_pretrained('''allenai/led-base-16384'''				).led

    # change to intended input here
    _UpperCamelCase:  List[str]    						=			_long_tensor([512 * [0, 31_414, 232, 328, 740, 1_140, 12_695, 69]]				)
    _UpperCamelCase:  int    						=			_long_tensor([128 * [0, 31_414, 232, 328, 740, 1_140, 12_695, 69]]				)
    _UpperCamelCase:  Union[str, Any]    						=			prepare_led_inputs_dict(model.config		,	_lowercase		,	_lowercase				)
    _UpperCamelCase:  Optional[Any]    						=			model(**_lowercase				)[0]
    _UpperCamelCase:  Optional[int]    						=			(1, 1_024, 768)
    self.assertEqual(output.shape		,	_lowercase				)
    # change to expected output here
    _UpperCamelCase:  Optional[Any]    						=			tf.convert_to_tensor(
        [[2.3050, 2.8279, 0.6531], [-1.8457, -0.1455, -3.5661], [-1.0186, 0.4586, -2.2043]]		,	)
    tf.debugging.assert_near(output[:, :3, :3]		,	_lowercase		,	atol=1E-3				)



   def 	lowerCAmelCase    (						self       :	Union[str, Any]				):



    """simple docstring"""



    _UpperCamelCase:  Optional[Any]    						=			TFLEDForConditionalGeneration.from_pretrained('''allenai/led-base-16384'''				)

    # change to intended input here
    _UpperCamelCase:  Dict    						=			_long_tensor([512 * [0, 31_414, 232, 328, 740, 1_140, 12_695, 69]]				)
    _UpperCamelCase:  List[str]    						=			_long_tensor([128 * [0, 31_414, 232, 328, 740, 1_140, 12_695, 69]]				)
    _UpperCamelCase:  Tuple    						=			prepare_led_inputs_dict(model.config		,	_lowercase		,	_lowercase				)
    _UpperCamelCase:  Any    						=			model(**_lowercase				)[0]
    _UpperCamelCase:  Dict    						=			(1, 1_024, model.config.vocab_size)
    self.assertEqual(output.shape		,	_lowercase				)
    # change to expected output here
    _UpperCamelCase:  Optional[Any]    						=			tf.convert_to_tensor(
        [[33.6507, 6.4572, 16.8089], [5.8739, -2.4238, 11.2902], [-3.2139, -4.3149, 4.2783]]		,	)
    tf.debugging.assert_near(output[:, :3, :3]		,	_lowercase		,	atol=1E-3		,	rtol=1E-3				)