



'''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

							__lowerCAmelCase									=   '''0.12'''  # assumed parallelism: 8


if is_torch_available():
							import torch



def      __lowerCamelCase     (      lowerCAmelCase_						,			lowerCAmelCase_						,			lowerCAmelCase_=None  )			->      Optional[Any]:
		if rng is None:
				_a     :      int        =							random.Random()

		_a     :      List[Any]        =							1
		for dim in shape:
				total_dims *= dim

		_a     :      Optional[int]        =							[]
		for _ in range(lowerCAmelCase_  ):
				values.append(rng.randint(0						,			vocab_size - 1  )  )

		_a     :      List[Any]        =							np.array(lowerCAmelCase_						,			dtype=jnp.intaa  ).reshape(lowerCAmelCase_  )

		return output



def      __lowerCamelCase     (      lowerCAmelCase_						,			lowerCAmelCase_=None  )			->      Union[str, Any]:
		_a     :      Union[str, Any]        =							ids_tensor(lowerCAmelCase_						,			vocab_size=2						,			rng=lowerCAmelCase_  )
		# make sure that at least one token is attended to for each batch
		_a     :      Dict        =							1
		return attn_mask

@require_flax
class 	__magic_name__						:
							lowerCAmelCase       :    List[Any]   				=					None
							lowerCAmelCase       :    Any   				=					()







							def 							__lowercase			(							self					:       List[Any]		):
									_a       ,					_a     :      Optional[int]        =							self.model_tester.prepare_config_and_inputs_for_common()

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

									_a     :      Any        =							jnp.ones_like(_UpperCAmelCase		)
									_a     :      Tuple        =							attention_mask[:max_batch_size, :sequence_length]

									# generate max 5 tokens
									_a     :      Optional[Any]        =							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     :      List[Any]        =							config.eos_token_id
									return config, input_ids, attention_mask, max_length







							@is_pt_flax_cross_test
							def 							__lowercase			(							self					:       Tuple		):
									_a       ,					_a       ,					_a       ,					_a     :      List[Any]        =							self._get_input_ids_and_config()
									_a     :      str        =							False
									_a     :      Dict        =							max_length
									_a     :      Union[str, Any]        =							0

									for model_class in self.all_generative_model_classes:
											_a     :      str        =							model_class(_UpperCAmelCase		)

											_a     :      Dict        =							model_class.__name__[4:]  # Skip the "Flax" at the beginning
											_a     :      List[str]        =							getattr(_UpperCAmelCase	,_UpperCAmelCase		)
											_a     :      List[str]        =							pt_model_class(_UpperCAmelCase		).eval()
											_a     :      Any        =							load_flax_weights_in_pytorch_model(_UpperCAmelCase	,flax_model.params		)

											_a     :      Optional[int]        =							flax_model.generate(_UpperCAmelCase		).sequences
											_a     :      Optional[int]        =							pt_model.generate(torch.tensor(_UpperCAmelCase	,dtype=torch.long		)		)

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

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







							def 							__lowercase			(							self					:       Any		):
									_a       ,					_a       ,					_a       ,					_a     :      Union[str, Any]        =							self._get_input_ids_and_config()
									_a     :      Tuple        =							False
									_a     :      Union[str, Any]        =							max_length

									for model_class in self.all_generative_model_classes:
											_a     :      Tuple        =							model_class(_UpperCAmelCase		)

											_a     :      str        =							model.generate(_UpperCAmelCase		).sequences
											self.assertEqual(generation_outputs.shape[-1]	,_UpperCAmelCase		)

											_a     :      int        =							jit(model.generate		)
											_a     :      str        =							jit_generate(_UpperCAmelCase		).sequences

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







							def 							__lowercase			(							self					:       List[str]		):
									_a       ,					_a       ,					_a       ,					_a     :      Union[str, Any]        =							self._get_input_ids_and_config()
									_a     :      str        =							True
									_a     :      Union[str, Any]        =							max_length

									for model_class in self.all_generative_model_classes:
											_a     :      Optional[Any]        =							model_class(_UpperCAmelCase		)

											_a     :      Union[str, Any]        =							model.generate(_UpperCAmelCase		).sequences
											self.assertEqual(generation_outputs.shape[-1]	,_UpperCAmelCase		)

											_a     :      List[str]        =							jit(model.generate		)
											_a     :      Union[str, Any]        =							jit_generate(_UpperCAmelCase		).sequences

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







							def 							__lowercase			(							self					:       Tuple		):
									_a       ,					_a       ,					_a       ,					_a     :      Optional[int]        =							self._get_input_ids_and_config()
									_a     :      Dict        =							False
									_a     :      Optional[int]        =							max_length
									_a     :      Optional[int]        =							2

									for model_class in self.all_generative_model_classes:
											_a     :      Dict        =							model_class(_UpperCAmelCase		)

											_a     :      Tuple        =							model.generate(_UpperCAmelCase		).sequences
											self.assertEqual(generation_outputs.shape[-1]	,_UpperCAmelCase		)

											_a     :      Optional[int]        =							jit(model.generate		)
											_a     :      Tuple        =							jit_generate(_UpperCAmelCase		).sequences

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







							def 							__lowercase			(							self					:       Optional[Any]		):
									_a       ,					_a       ,					_a       ,					_a     :      Dict        =							self._get_input_ids_and_config()
									_a     :      Tuple        =							False
									_a     :      Dict        =							max_length
									_a     :      int        =							2
									_a     :      Dict        =							2

									for model_class in self.all_generative_model_classes:
											_a     :      Any        =							model_class(_UpperCAmelCase		)

											_a     :      List[str]        =							model.generate(_UpperCAmelCase		).sequences
											self.assertEqual(generation_outputs.shape[0]	,input_ids.shape[0] * config.num_return_sequences		)







							def 							__lowercase			(							self					:       str		):
									_a       ,					_a       ,					_a       ,					_a     :      Optional[Any]        =							self._get_input_ids_and_config()
									_a     :      List[str]        =							True
									_a     :      Tuple        =							max_length
									_a     :      int        =							0.8
									_a     :      List[Any]        =							10
									_a     :      List[Any]        =							0.3
									_a     :      Optional[int]        =							1
									_a     :      Union[str, Any]        =							8
									_a     :      int        =							9

									for model_class in self.all_generative_model_classes:
											_a     :      Optional[Any]        =							model_class(_UpperCAmelCase		)

											_a     :      str        =							model.generate(_UpperCAmelCase		).sequences
											self.assertEqual(generation_outputs.shape[-1]	,_UpperCAmelCase		)

											_a     :      List[str]        =							jit(model.generate		)
											_a     :      Dict        =							jit_generate(_UpperCAmelCase		).sequences

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







							def 							__lowercase			(							self					:       int		):
									_a       ,					_a       ,					_a       ,					_a     :      Any        =							self._get_input_ids_and_config()
									_a     :      Union[str, Any]        =							max_length
									_a     :      Optional[Any]        =							1
									_a     :      List[Any]        =							8
									_a     :      Optional[int]        =							9

									for model_class in self.all_generative_model_classes:
											_a     :      Tuple        =							model_class(_UpperCAmelCase		)

											_a     :      Any        =							model.generate(_UpperCAmelCase		).sequences
											self.assertEqual(generation_outputs.shape[-1]	,_UpperCAmelCase		)

											_a     :      List[str]        =							jit(model.generate		)
											_a     :      Any        =							jit_generate(_UpperCAmelCase		).sequences

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







							def 							__lowercase			(							self					:       Optional[int]		):
									_a       ,					_a       ,					_a       ,					_a     :      Tuple        =							self._get_input_ids_and_config()
									_a     :      Tuple        =							max_length
									_a     :      Any        =							2
									_a     :      Tuple        =							1
									_a     :      Any        =							8
									_a     :      Optional[int]        =							9

									for model_class in self.all_generative_model_classes:
											_a     :      str        =							model_class(_UpperCAmelCase		)

											_a     :      Tuple        =							model.generate(_UpperCAmelCase		).sequences
											self.assertEqual(generation_outputs.shape[-1]	,_UpperCAmelCase		)

											_a     :      Tuple        =							jit(model.generate		)
											_a     :      List[Any]        =							jit_generate(_UpperCAmelCase		).sequences

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







							def 							__lowercase			(							self					:       Union[str, Any]		):
									_a       ,					_a       ,					_a       ,					_a     :      Union[str, Any]        =							self._get_input_ids_and_config()

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

									_a     :      List[Any]        =							False
									_a     :      Optional[Any]        =							max_length

									for model_class in self.all_generative_model_classes:
											_a     :      int        =							model_class(_UpperCAmelCase		)

											_a     :      str        =							model.generate(_UpperCAmelCase	,attention_mask=_UpperCAmelCase		).sequences
											self.assertEqual(generation_outputs.shape[-1]	,_UpperCAmelCase		)

											_a     :      List[Any]        =							jit(model.generate		)
											_a     :      List[Any]        =							jit_generate(_UpperCAmelCase	,attention_mask=_UpperCAmelCase		).sequences

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







							def 							__lowercase			(							self					:       Optional[Any]		):
									_a       ,					_a       ,					_a       ,					_a     :      List[Any]        =							self._get_input_ids_and_config()

									# pad attention mask on the left
									_a     :      Optional[Any]        =							attention_mask.at[(0, 0)].set(0		)

									_a     :      Dict        =							True
									_a     :      Union[str, Any]        =							max_length

									for model_class in self.all_generative_model_classes:
											_a     :      int        =							model_class(_UpperCAmelCase		)

											_a     :      Dict        =							model.generate(_UpperCAmelCase	,attention_mask=_UpperCAmelCase		).sequences
											self.assertEqual(generation_outputs.shape[-1]	,_UpperCAmelCase		)

											_a     :      Optional[int]        =							jit(model.generate		)
											_a     :      Tuple        =							jit_generate(_UpperCAmelCase	,attention_mask=_UpperCAmelCase		).sequences

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







							def 							__lowercase			(							self					:       Optional[int]		):
									_a       ,					_a       ,					_a       ,					_a     :      Any        =							self._get_input_ids_and_config()

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

									_a     :      Dict        =							2
									_a     :      List[Any]        =							max_length

									for model_class in self.all_generative_model_classes:
											_a     :      str        =							model_class(_UpperCAmelCase		)

											_a     :      Optional[Any]        =							model.generate(_UpperCAmelCase	,attention_mask=_UpperCAmelCase		).sequences
											self.assertEqual(generation_outputs.shape[-1]	,_UpperCAmelCase		)

											_a     :      Optional[Any]        =							jit(model.generate		)
											_a     :      List[str]        =							jit_generate(_UpperCAmelCase	,attention_mask=_UpperCAmelCase		).sequences

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

@require_flax
class 	__magic_name__						(    unittest.TestCase   ):







							def 							__lowercase			(							self					:       Union[str, Any]		):
									_a     :      Any        =							AutoTokenizer.from_pretrained('hf-internal-testing/tiny-bert'		)
									_a     :      Optional[int]        =							FlaxAutoModelForCausalLM.from_pretrained('hf-internal-testing/tiny-bert-flax-only'		)

									_a     :      Optional[int]        =							'Hello world'
									_a     :      Optional[Any]        =							tokenizer(_UpperCAmelCase	,return_tensors='np'		).input_ids

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

									# arbitrary arguments that will not be used anywhere are also not accepted
									with self.assertRaisesRegex(_UpperCAmelCase	,'foo'		):
											_a     :      Optional[int]        =							{'foo': 'bar'}
											model.generate(_UpperCAmelCase	,**_UpperCAmelCase		)












"""simple docstring"""






def 					_A			(__a = 50							) ->       int:
    """simple docstring"""






    SCREAMING_SNAKE_CASE_      :  Tuple    						=		[[0] * 3 for _ in range(length + 1							)]

    for row_length in range(length + 1							):
        for tile_length in range(2    ,				5							):
            for tile_start in range(row_length - tile_length + 1							):
                different_colour_ways_number[row_length][tile_length - 2] += (
                    different_colour_ways_number[row_length - tile_start - tile_length][
                        tile_length - 2
                    ]
                    + 1
                )

    return sum(different_colour_ways_number[length]							)


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


