



"""simple docstring"""







import random
import unittest

import numpy as np

import transformers
from transformers import is_flax_available, is_torch_available
from transformers.testing_utils import is_pt_flax_cross_test, require_flax


if is_flax_available():
			import os

			import jax.numpy as jnp
			from jax import jit

			from transformers import AutoTokenizer, FlaxAutoModelForCausalLM
			from transformers.modeling_flax_pytorch_utils import load_flax_weights_in_pytorch_model

			UpperCamelCase     =							"""0.12"""  # assumed parallelism: 8


if is_torch_available():
			import torch

def       _lowerCamelCase				(							UpperCAmelCase_ :       Union[str, Any],	UpperCAmelCase_ :       List[str],	UpperCAmelCase_ :       str=None     )	-> List[Any]:






	"""simple docstring"""





	if rng is None:
		A__											=					random.Random()

	A__											=					1
	for dim in shape:
		total_dims *= dim

	A__											=					[]
	for _ in range(UpperCAmelCase_     ):
		values.append(rng.randint(0,	vocab_size - 1     )     )

	A__											=					np.array(UpperCAmelCase_,	dtype=jnp.intaa     ).reshape(UpperCAmelCase_     )

	return output

def       _lowerCamelCase				(							UpperCAmelCase_ :       str,	UpperCAmelCase_ :       Optional[int]=None     )	-> str:






	"""simple docstring"""





	A__											=					ids_tensor(UpperCAmelCase_,	vocab_size=2,	rng=UpperCAmelCase_     )
	# make sure that at least one token is attended to for each batch
	A__											=					1
	return attn_mask






@require_flax
class       UpperCamelCase__			:

			"""simple docstring"""


			A__	:					Union[str, Any]							  =     None
			A__	:					Optional[Any]							  =     ()


			def   snake_case__	(      self		)      ->	Dict:
				A__											=					self.model_tester.prepare_config_and_inputs_for_common()

				# cut to half length & take max batch_size 3
				A__											=					2
				A__											=					inputs['input_ids'].shape[-1] // 2
				A__											=					inputs['input_ids'][:max_batch_size, :sequence_length]

				A__											=					jnp.ones_like(SCREAMING_SNAKE_CASE_		)
				A__											=					attention_mask[:max_batch_size, :sequence_length]

				# generate max 5 tokens
				A__											=					input_ids.shape[-1] + 5
				if config.eos_token_id is not None and config.pad_token_id is None:
					# hack to allow generate for models such as GPT2 as is done in `generate()`
					A__											=					config.eos_token_id
				return config, input_ids, attention_mask, max_length


			@is_pt_flax_cross_test
			def   snake_case__	(      self		)      ->	List[Any]:
				A__											=					self._get_input_ids_and_config()
				A__											=					False
				A__											=					max_length
				A__											=					0

				for model_class in self.all_generative_model_classes:
					A__											=					model_class(SCREAMING_SNAKE_CASE_		)

					A__											=					model_class.__name__[4:]  # Skip the "Flax" at the beginning
					A__											=					getattr(SCREAMING_SNAKE_CASE_						,	SCREAMING_SNAKE_CASE_		)
					A__											=					pt_model_class(SCREAMING_SNAKE_CASE_		).eval()
					A__											=					load_flax_weights_in_pytorch_model(SCREAMING_SNAKE_CASE_						,	flax_model.params		)

					A__											=					flax_model.generate(SCREAMING_SNAKE_CASE_		).sequences
					A__											=					pt_model.generate(torch.tensor(SCREAMING_SNAKE_CASE_						,	dtype=torch.long		)		)

					if flax_generation_outputs.shape[-1] > pt_generation_outputs.shape[-1]:
						A__											=					flax_generation_outputs[:, : pt_generation_outputs.shape[-1]]

					self.assertListEqual(pt_generation_outputs.numpy().tolist()						,	flax_generation_outputs.tolist()		)


			def   snake_case__	(      self		)      ->	str:
				A__											=					self._get_input_ids_and_config()
				A__											=					False
				A__											=					max_length

				for model_class in self.all_generative_model_classes:
					A__											=					model_class(SCREAMING_SNAKE_CASE_		)

					A__											=					model.generate(SCREAMING_SNAKE_CASE_		).sequences
					self.assertEqual(generation_outputs.shape[-1]						,	SCREAMING_SNAKE_CASE_		)

					A__											=					jit(model.generate		)
					A__											=					jit_generate(SCREAMING_SNAKE_CASE_		).sequences

					self.assertListEqual(generation_outputs.tolist()						,	jit_generation_outputs.tolist()		)


			def   snake_case__	(      self		)      ->	Union[str, Any]:
				A__											=					self._get_input_ids_and_config()
				A__											=					True
				A__											=					max_length

				for model_class in self.all_generative_model_classes:
					A__											=					model_class(SCREAMING_SNAKE_CASE_		)

					A__											=					model.generate(SCREAMING_SNAKE_CASE_		).sequences
					self.assertEqual(generation_outputs.shape[-1]						,	SCREAMING_SNAKE_CASE_		)

					A__											=					jit(model.generate		)
					A__											=					jit_generate(SCREAMING_SNAKE_CASE_		).sequences

					self.assertListEqual(generation_outputs.tolist()						,	jit_generation_outputs.tolist()		)


			def   snake_case__	(      self		)      ->	List[Any]:
				A__											=					self._get_input_ids_and_config()
				A__											=					False
				A__											=					max_length
				A__											=					2

				for model_class in self.all_generative_model_classes:
					A__											=					model_class(SCREAMING_SNAKE_CASE_		)

					A__											=					model.generate(SCREAMING_SNAKE_CASE_		).sequences
					self.assertEqual(generation_outputs.shape[-1]						,	SCREAMING_SNAKE_CASE_		)

					A__											=					jit(model.generate		)
					A__											=					jit_generate(SCREAMING_SNAKE_CASE_		).sequences

					self.assertListEqual(generation_outputs.tolist()						,	jit_generation_outputs.tolist()		)


			def   snake_case__	(      self		)      ->	Optional[int]:
				A__											=					self._get_input_ids_and_config()
				A__											=					False
				A__											=					max_length
				A__											=					2
				A__											=					2

				for model_class in self.all_generative_model_classes:
					A__											=					model_class(SCREAMING_SNAKE_CASE_		)

					A__											=					model.generate(SCREAMING_SNAKE_CASE_		).sequences
					self.assertEqual(generation_outputs.shape[0]						,	input_ids.shape[0] * config.num_return_sequences		)


			def   snake_case__	(      self		)      ->	Dict:
				A__											=					self._get_input_ids_and_config()
				A__											=					True
				A__											=					max_length
				A__											=					0.8
				A__											=					10
				A__											=					0.3
				A__											=					1
				A__											=					8
				A__											=					9

				for model_class in self.all_generative_model_classes:
					A__											=					model_class(SCREAMING_SNAKE_CASE_		)

					A__											=					model.generate(SCREAMING_SNAKE_CASE_		).sequences
					self.assertEqual(generation_outputs.shape[-1]						,	SCREAMING_SNAKE_CASE_		)

					A__											=					jit(model.generate		)
					A__											=					jit_generate(SCREAMING_SNAKE_CASE_		).sequences

					self.assertListEqual(generation_outputs.tolist()						,	jit_generation_outputs.tolist()		)


			def   snake_case__	(      self		)      ->	Optional[int]:
				A__											=					self._get_input_ids_and_config()
				A__											=					max_length
				A__											=					1
				A__											=					8
				A__											=					9

				for model_class in self.all_generative_model_classes:
					A__											=					model_class(SCREAMING_SNAKE_CASE_		)

					A__											=					model.generate(SCREAMING_SNAKE_CASE_		).sequences
					self.assertEqual(generation_outputs.shape[-1]						,	SCREAMING_SNAKE_CASE_		)

					A__											=					jit(model.generate		)
					A__											=					jit_generate(SCREAMING_SNAKE_CASE_		).sequences

					self.assertListEqual(generation_outputs.tolist()						,	jit_generation_outputs.tolist()		)


			def   snake_case__	(      self		)      ->	Optional[int]:
				A__											=					self._get_input_ids_and_config()
				A__											=					max_length
				A__											=					2
				A__											=					1
				A__											=					8
				A__											=					9

				for model_class in self.all_generative_model_classes:
					A__											=					model_class(SCREAMING_SNAKE_CASE_		)

					A__											=					model.generate(SCREAMING_SNAKE_CASE_		).sequences
					self.assertEqual(generation_outputs.shape[-1]						,	SCREAMING_SNAKE_CASE_		)

					A__											=					jit(model.generate		)
					A__											=					jit_generate(SCREAMING_SNAKE_CASE_		).sequences

					self.assertListEqual(generation_outputs.tolist()						,	jit_generation_outputs.tolist()		)


			def   snake_case__	(      self		)      ->	str:
				A__											=					self._get_input_ids_and_config()

				# pad attention mask on the left
				A__											=					attention_mask.at[(0, 0)].set(0		)

				A__											=					False
				A__											=					max_length

				for model_class in self.all_generative_model_classes:
					A__											=					model_class(SCREAMING_SNAKE_CASE_		)

					A__											=					model.generate(SCREAMING_SNAKE_CASE_						,	attention_mask=SCREAMING_SNAKE_CASE_		).sequences
					self.assertEqual(generation_outputs.shape[-1]						,	SCREAMING_SNAKE_CASE_		)

					A__											=					jit(model.generate		)
					A__											=					jit_generate(SCREAMING_SNAKE_CASE_						,	attention_mask=SCREAMING_SNAKE_CASE_		).sequences

					self.assertListEqual(generation_outputs.tolist()						,	jit_generation_outputs.tolist()		)


			def   snake_case__	(      self		)      ->	List[Any]:
				A__											=					self._get_input_ids_and_config()

				# pad attention mask on the left
				A__											=					attention_mask.at[(0, 0)].set(0		)

				A__											=					True
				A__											=					max_length

				for model_class in self.all_generative_model_classes:
					A__											=					model_class(SCREAMING_SNAKE_CASE_		)

					A__											=					model.generate(SCREAMING_SNAKE_CASE_						,	attention_mask=SCREAMING_SNAKE_CASE_		).sequences
					self.assertEqual(generation_outputs.shape[-1]						,	SCREAMING_SNAKE_CASE_		)

					A__											=					jit(model.generate		)
					A__											=					jit_generate(SCREAMING_SNAKE_CASE_						,	attention_mask=SCREAMING_SNAKE_CASE_		).sequences

					self.assertListEqual(generation_outputs.tolist()						,	jit_generation_outputs.tolist()		)







			def   snake_case__	(      self		)      ->	Optional[int]:
				A__											=					self._get_input_ids_and_config()

				# pad attention mask on the left
				A__											=					attention_mask.at[(0, 0)].set(0		)

				A__											=					2
				A__											=					max_length

				for model_class in self.all_generative_model_classes:
					A__											=					model_class(SCREAMING_SNAKE_CASE_		)

					A__											=					model.generate(SCREAMING_SNAKE_CASE_						,	attention_mask=SCREAMING_SNAKE_CASE_		).sequences
					self.assertEqual(generation_outputs.shape[-1]						,	SCREAMING_SNAKE_CASE_		)

					A__											=					jit(model.generate		)
					A__											=					jit_generate(SCREAMING_SNAKE_CASE_						,	attention_mask=SCREAMING_SNAKE_CASE_		).sequences

					self.assertListEqual(generation_outputs.tolist()						,	jit_generation_outputs.tolist()		)






@require_flax
class       UpperCamelCase__			(   unittest.TestCase      ):

			"""simple docstring"""


			def   snake_case__	(      self		)      ->	List[Any]:
				A__											=					AutoTokenizer.from_pretrained("hf-internal-testing/tiny-bert"		)
				A__											=					FlaxAutoModelForCausalLM.from_pretrained("hf-internal-testing/tiny-bert-flax-only"		)

				A__											=					'Hello world'
				A__											=					tokenizer(SCREAMING_SNAKE_CASE_						,	return_tensors="np"		).input_ids

				# typos are quickly detected (the correct argument is `do_sample`)
				with self.assertRaisesRegex(SCREAMING_SNAKE_CASE_						,	"do_samples"		):
					model.generate(SCREAMING_SNAKE_CASE_						,	do_samples=SCREAMING_SNAKE_CASE_		)

				# arbitrary arguments that will not be used anywhere are also not accepted
				with self.assertRaisesRegex(SCREAMING_SNAKE_CASE_						,	"foo"		):
					A__											=					{'foo': 'bar'}
					model.generate(SCREAMING_SNAKE_CASE_						,	**SCREAMING_SNAKE_CASE_		)








'''simple docstring'''



A__						:   dict[tuple[int, int, int], int]				=    {}



def      UpperCAmelCase__   (      UpperCAmelCase_       :							int ,       UpperCAmelCase_       :							int ,       UpperCAmelCase_       :							int				)	->						int:

     # if we are absent twice, or late 3 consecutive days,
     # no further prize strings are possible
     if late == 3 or absent == 2:
          return 0

     # if we have no days left, and have not failed any other rules,
     # we have a prize string
     if days == 0:
          return 1

     # No easy solution, so now we need to do the recursive calculation

     # First, check if the combination is already in the cache, and
     # if yes, return the stored value from there since we already
     # know the number of possible prize strings from this point on
     __lowerCamelCase  :      List[Any]														=						(days, absent, late)
     if key in cache:
          return cache[key]

     # now we calculate the three possible ways that can unfold from
     # this point on, depending on our attendance today

     # 1) if we are late (but not absent), the "absent" counter stays as
     # it is, but the "late" counter increases by one
     __lowerCamelCase  :      Tuple														=						_calculate(days - 1 ,       UpperCAmelCase_ ,       late + 1				)

     # 2) if we are absent, the "absent" counter increases by 1, and the
     # "late" counter resets to 0
     __lowerCamelCase  :      int														=						_calculate(days - 1 ,       absent + 1 ,       0				)

     # 3) if we are on time, this resets the "late" counter and keeps the
     # absent counter
     __lowerCamelCase  :      List[Any]														=						_calculate(days - 1 ,       UpperCAmelCase_ ,       0				)

     __lowerCamelCase  :      Optional[int]														=						state_late + state_absent + state_ontime

     __lowerCamelCase  :      Union[str, Any]														=						prizestrings
     return prizestrings



def      UpperCAmelCase__   (      UpperCAmelCase_       :							int = 30				)	->						int:
     return _calculate(UpperCAmelCase_ ,       absent=0 ,       late=0				)


if __name__ == "__main__":
      print(solution())






