



"""simple docstring"""
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       SCREAMING_SNAKE_CASE__  :
   __lowerCAmelCase  :   Optional[Any]						=						LEDConfig
   __lowerCAmelCase  :   Any						=						{}
   __lowerCAmelCase  :   Optional[Any]						=						'gelu'

   def __init__(   self      ,     _SCREAMING_SNAKE_CASE      ,     _SCREAMING_SNAKE_CASE=13      ,     _SCREAMING_SNAKE_CASE=7      ,     _SCREAMING_SNAKE_CASE=True      ,     _SCREAMING_SNAKE_CASE=False      ,     _SCREAMING_SNAKE_CASE=99      ,     _SCREAMING_SNAKE_CASE=32      ,     _SCREAMING_SNAKE_CASE=2      ,     _SCREAMING_SNAKE_CASE=4      ,     _SCREAMING_SNAKE_CASE=37      ,     _SCREAMING_SNAKE_CASE=0.1      ,     _SCREAMING_SNAKE_CASE=0.1      ,     _SCREAMING_SNAKE_CASE=20      ,     _SCREAMING_SNAKE_CASE=2      ,     _SCREAMING_SNAKE_CASE=1      ,     _SCREAMING_SNAKE_CASE=0      ,     _SCREAMING_SNAKE_CASE=4      ,     )  ->      List[Any]:





    '''simple docstring'''




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

    UpperCAmelCase	:					List[str]	=	hidden_dropout_prob
    UpperCAmelCase	:					List[str]	=	attention_probs_dropout_prob
    UpperCAmelCase	:					List[Any]	=	max_position_embeddings
    UpperCAmelCase	:					Optional[int]	=	eos_token_id
    UpperCAmelCase	:					str	=	pad_token_id
    UpperCAmelCase	:					Tuple	=	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	:					int	=	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	:					str	=	(
        self.seq_length + (self.attention_window - self.seq_length % self.attention_window) % self.attention_window
    )

   def 	SCREAMING_SNAKE_CASE     (   self						)  ->      Any:





    '''simple docstring'''




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

    UpperCAmelCase	:					Optional[Any]	=	ids_tensor([self.batch_size, self.seq_length]      ,     self.vocab_size						)

    UpperCAmelCase	:					Tuple	=	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	:					Optional[int]	=	prepare_led_inputs_dict(_SCREAMING_SNAKE_CASE      ,     _SCREAMING_SNAKE_CASE      ,     _SCREAMING_SNAKE_CASE						)
    UpperCAmelCase	:					Any	=	tf.concat(
        [tf.zeros_like(_SCREAMING_SNAKE_CASE						)[:, :-1], tf.ones_like(_SCREAMING_SNAKE_CASE						)[:, -1:]]      ,     axis=-1      ,     )
    UpperCAmelCase	:					Union[str, Any]	=	global_attention_mask
    return config, inputs_dict




   def 	SCREAMING_SNAKE_CASE     (   self      ,     _SCREAMING_SNAKE_CASE      ,     _SCREAMING_SNAKE_CASE						)  ->      int:





    '''simple docstring'''




    UpperCAmelCase	:					Tuple	=	TFLEDModel(config=_SCREAMING_SNAKE_CASE						).get_decoder()
    UpperCAmelCase	:					Optional[Any]	=	inputs_dict["""input_ids"""]

    UpperCAmelCase	:					Any	=	input_ids[:1, :]
    UpperCAmelCase	:					int	=	inputs_dict["""attention_mask"""][:1, :]
    UpperCAmelCase	:					Optional[Any]	=	1

    # first forward pass
    UpperCAmelCase	:					Optional[int]	=	model(_SCREAMING_SNAKE_CASE      ,     attention_mask=_SCREAMING_SNAKE_CASE      ,     use_cache=_SCREAMING_SNAKE_CASE						)

    UpperCAmelCase	,		UpperCAmelCase	:					List[Any]	=	outputs.to_tuple()

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

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

    UpperCAmelCase	:					List[Any]	=	model(_SCREAMING_SNAKE_CASE      ,     attention_mask=_SCREAMING_SNAKE_CASE						)[0]
    UpperCAmelCase	:					Union[str, Any]	=	model(_SCREAMING_SNAKE_CASE      ,     attention_mask=_SCREAMING_SNAKE_CASE      ,     past_key_values=_SCREAMING_SNAKE_CASE						)[0]

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

    # select random slice
    UpperCAmelCase	:					int	=	int(ids_tensor((1,)      ,     output_from_past.shape[-1]						)						)
    UpperCAmelCase	:					int	=	output_from_no_past[:, -3:, random_slice_idx]
    UpperCAmelCase	:					str	=	output_from_past[:, :, random_slice_idx]

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

def       _snake_case						(						UpperCamelCase						:		Any			,		UpperCamelCase						:		int			,		UpperCamelCase						:		Union[str, Any]			,		UpperCamelCase						:		List[str]=None			,		UpperCamelCase						:		str=None			,		UpperCamelCase						:		Dict=None			,		UpperCamelCase						:		List[Any]=None			,		):
 if attention_mask is None:
  UpperCAmelCase	:					int	=	tf.cast(tf.math.not_equal(UpperCamelCase			,		config.pad_token_id						)			,		tf.inta						)
 if decoder_attention_mask is None:
  UpperCAmelCase	:					Dict	=	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	:					List[Any]	=	tf.ones((config.encoder_layers, config.encoder_attention_heads)						)
 if decoder_head_mask is None:
  UpperCAmelCase	:					Union[str, 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       SCREAMING_SNAKE_CASE__  (					UpperCAmelCase__      ,			UpperCAmelCase__      ,			unittest.TestCase ):
   __lowerCAmelCase  :   Optional[Any]						=						(TFLEDForConditionalGeneration, TFLEDModel) if is_tf_available() else ()
   __lowerCAmelCase  :   Dict						=						(TFLEDForConditionalGeneration,) if is_tf_available() else ()
   __lowerCAmelCase  :   int						=						(
       {
           'conversational': TFLEDForConditionalGeneration,
           'feature-extraction': TFLEDModel,
           'summarization': TFLEDForConditionalGeneration,
           'text2text-generation': TFLEDForConditionalGeneration,
           'translation': TFLEDForConditionalGeneration,
       }
       if is_tf_available()
       else {}
   )
   __lowerCAmelCase  :   Any						=						True
   __lowerCAmelCase  :   Tuple						=						False
   __lowerCAmelCase  :   Union[str, Any]						=						False
   __lowerCAmelCase  :   Dict						=						False

   def 	SCREAMING_SNAKE_CASE     (   self						)  ->      str:





    '''simple docstring'''




    UpperCAmelCase	:					Optional[int]	=	TFLEDModelTester(self						)
    UpperCAmelCase	:					int	=	ConfigTester(self      ,     config_class=_SCREAMING_SNAKE_CASE						)

   def 	SCREAMING_SNAKE_CASE     (   self						)  ->      Optional[Any]:





    '''simple docstring'''




    self.config_tester.run_common_tests()

   def 	SCREAMING_SNAKE_CASE     (   self						)  ->      Any:





    '''simple docstring'''




    UpperCAmelCase	:					Dict	=	self.model_tester.prepare_config_and_inputs_for_common()
    self.model_tester.check_decoder_model_past_large_inputs(*_SCREAMING_SNAKE_CASE						)

   def 	SCREAMING_SNAKE_CASE     (   self						)  ->      Dict:





    '''simple docstring'''




    UpperCAmelCase	,		UpperCAmelCase	:					int	=	self.model_tester.prepare_config_and_inputs_for_common()
    UpperCAmelCase	:					Union[str, Any]	=	tf.zeros_like(inputs_dict["""attention_mask"""]						)
    UpperCAmelCase	:					Tuple	=	2
    UpperCAmelCase	:					Tuple	=	tf.where(
        tf.range(self.model_tester.seq_length						)[None, :] < num_global_attn_indices      ,     1      ,     inputs_dict["""global_attention_mask"""]      ,     )

    UpperCAmelCase	:					List[Any]	=	True
    UpperCAmelCase	:					int	=	self.model_tester.seq_length
    UpperCAmelCase	:					Optional[int]	=	self.model_tester.encoder_seq_length

    def check_decoder_attentions_output(_SCREAMING_SNAKE_CASE						):
     UpperCAmelCase	:					str	=	outputs.decoder_attentions
     self.assertEqual(len(_SCREAMING_SNAKE_CASE						)      ,     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(_SCREAMING_SNAKE_CASE						):
     UpperCAmelCase	:					List[Any]	=	[t.numpy() for t in outputs.encoder_attentions]
     UpperCAmelCase	:					int	=	[t.numpy() for t in outputs.encoder_global_attentions]
     self.assertEqual(len(_SCREAMING_SNAKE_CASE						)      ,     self.model_tester.num_hidden_layers						)
     self.assertEqual(len(_SCREAMING_SNAKE_CASE						)      ,     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[Any]	=	True
     UpperCAmelCase	:					List[str]	=	False
     UpperCAmelCase	:					str	=	False
     UpperCAmelCase	:					Any	=	model_class(_SCREAMING_SNAKE_CASE						)
     UpperCAmelCase	:					Union[str, Any]	=	model(self._prepare_for_class(_SCREAMING_SNAKE_CASE      ,     _SCREAMING_SNAKE_CASE						)						)
     UpperCAmelCase	:					Tuple	=	len(_SCREAMING_SNAKE_CASE						)
     self.assertEqual(config.output_hidden_states      ,     _SCREAMING_SNAKE_CASE						)
     check_encoder_attentions_output(_SCREAMING_SNAKE_CASE						)

     if self.is_encoder_decoder:
      UpperCAmelCase	:					Optional[int]	=	model_class(_SCREAMING_SNAKE_CASE						)
      UpperCAmelCase	:					List[str]	=	model(self._prepare_for_class(_SCREAMING_SNAKE_CASE      ,     _SCREAMING_SNAKE_CASE						)						)
      self.assertEqual(config.output_hidden_states      ,     _SCREAMING_SNAKE_CASE						)
      check_decoder_attentions_output(_SCREAMING_SNAKE_CASE						)

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

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

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

   @unittest.skip("""LED keeps using potentially symbolic tensors in conditionals and breaks tracing."""						)
   def 	SCREAMING_SNAKE_CASE     (   self						)  ->      Any:





    '''simple docstring'''




    pass




   def 	SCREAMING_SNAKE_CASE     (   self						)  ->      List[str]:





    '''simple docstring'''




    pass

def       _snake_case						(						UpperCamelCase						:		Optional[Any]						):
 return tf.constant(UpperCamelCase			,		dtype=tf.intaa						)


A:   Tuple								=		1E-4



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

   def 	SCREAMING_SNAKE_CASE     (   self						)  ->      List[Any]:





    '''simple docstring'''




    UpperCAmelCase	:					Optional[int]	=	TFLEDForConditionalGeneration.from_pretrained("""allenai/led-base-16384"""						).led

    # change to intended input here
    UpperCAmelCase	:					int	=	_long_tensor([512 * [0, 31414, 232, 328, 740, 1140, 12695, 69]]						)
    UpperCAmelCase	:					List[Any]	=	_long_tensor([128 * [0, 31414, 232, 328, 740, 1140, 12695, 69]]						)
    UpperCAmelCase	:					Union[str, Any]	=	prepare_led_inputs_dict(model.config      ,     _SCREAMING_SNAKE_CASE      ,     _SCREAMING_SNAKE_CASE						)
    UpperCAmelCase	:					str	=	model(**_SCREAMING_SNAKE_CASE						)[0]
    UpperCAmelCase	:					List[str]	=	(1, 1024, 768)
    self.assertEqual(output.shape      ,     _SCREAMING_SNAKE_CASE						)
    # change to expected output here
    UpperCAmelCase	:					Dict	=	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]      ,     _SCREAMING_SNAKE_CASE      ,     atol=1E-3						)




   def 	SCREAMING_SNAKE_CASE     (   self						)  ->      List[Any]:





    '''simple docstring'''




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

    # change to intended input here
    UpperCAmelCase	:					Optional[Any]	=	_long_tensor([512 * [0, 31414, 232, 328, 740, 1140, 12695, 69]]						)
    UpperCAmelCase	:					List[str]	=	_long_tensor([128 * [0, 31414, 232, 328, 740, 1140, 12695, 69]]						)
    UpperCAmelCase	:					Optional[Any]	=	prepare_led_inputs_dict(model.config      ,     _SCREAMING_SNAKE_CASE      ,     _SCREAMING_SNAKE_CASE						)
    UpperCAmelCase	:					int	=	model(**_SCREAMING_SNAKE_CASE						)[0]
    UpperCAmelCase	:					List[Any]	=	(1, 1024, model.config.vocab_size)
    self.assertEqual(output.shape      ,     _SCREAMING_SNAKE_CASE						)
    # change to expected output here
    UpperCAmelCase	:					Union[str, 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]      ,     _SCREAMING_SNAKE_CASE      ,     atol=1E-3      ,     rtol=1E-3						)





"""simple docstring"""
from collections import defaultdict
from math import ceil, sqrt

def       _snake_case						(						UpperCamelCase						:		int = 1000000			,		UpperCamelCase						:		int = 10						):
 UpperCAmelCase	:					defaultdict	=	defaultdict(UpperCamelCase						)

 for outer_width in range(3			,		(t_limit // 4) + 2						):
  if outer_width * outer_width > t_limit:
   UpperCAmelCase	:					str	=	max(
       ceil(sqrt(outer_width * outer_width - t_limit						)						)			,		1						)
  else:
   UpperCAmelCase	:					Optional[Any]	=	1

  hole_width_lower_bound += (outer_width - hole_width_lower_bound) % 2

  for hole_width in range(UpperCamelCase			,		outer_width - 1			,		2						):
   count[outer_width * outer_width - hole_width * hole_width] += 1

 return sum(1 for n in count.values() if 1 <= n <= 10						)


if __name__ == "__main__":
      print(f"""{solution() = }""")
