



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     :							Optional[int]  							=				"0.12"  # assumed parallelism: 8


if is_torch_available():
		import torch

def    lowerCamelCase_(lowerCamelCase_    ,   lowerCamelCase_    ,   lowerCamelCase_=None    ) ->       Optional[Any]:
						if rng is None:
												UpperCAmelCase         =     random.Random()

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

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

						UpperCAmelCase         =     np.array(lowerCamelCase_    ,   dtype=jnp.intaa    ).reshape(lowerCamelCase_    )

						return output

def    lowerCamelCase_(lowerCamelCase_    ,   lowerCamelCase_=None    ) ->       Dict:
						UpperCAmelCase         =     ids_tensor(lowerCamelCase_    ,   vocab_size=2    ,   rng=lowerCamelCase_    )
						# make sure that at least one token is attended to for each batch
						UpperCAmelCase         =     1
						return attn_mask





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






		def SCREAMING_SNAKE_CASE_			(						self      :     List[Any]      )				->					Tuple:
								'''simple docstring'''







								UpperCAmelCase       ,		UpperCAmelCase         =     self.model_tester.prepare_config_and_inputs_for_common()

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

								UpperCAmelCase         =     jnp.ones_like(UpperCamelCase__      )
								UpperCAmelCase         =     attention_mask[:max_batch_size, :sequence_length]

								# generate max 5 tokens
								UpperCAmelCase         =     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()`
														UpperCAmelCase         =     config.eos_token_id
								return config, input_ids, attention_mask, max_length






		@is_pt_flax_cross_test
		def SCREAMING_SNAKE_CASE_			(						self      :     str      )				->					int:
								'''simple docstring'''







								UpperCAmelCase       ,		UpperCAmelCase       ,		UpperCAmelCase       ,		UpperCAmelCase         =     self._get_input_ids_and_config()
								UpperCAmelCase         =     False
								UpperCAmelCase         =     max_length
								UpperCAmelCase         =     0

								for model_class in self.all_generative_model_classes:
														UpperCAmelCase         =     model_class(UpperCamelCase__      )

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

														UpperCAmelCase         =     flax_model.generate(UpperCamelCase__      ).sequences
														UpperCAmelCase         =     pt_model.generate(torch.tensor(UpperCamelCase__				,		dtype=torch.long      )      )

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

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






		def SCREAMING_SNAKE_CASE_			(						self      :     Tuple      )				->					Optional[int]:
								'''simple docstring'''







								UpperCAmelCase       ,		UpperCAmelCase       ,		UpperCAmelCase       ,		UpperCAmelCase         =     self._get_input_ids_and_config()
								UpperCAmelCase         =     False
								UpperCAmelCase         =     max_length

								for model_class in self.all_generative_model_classes:
														UpperCAmelCase         =     model_class(UpperCamelCase__      )

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

														UpperCAmelCase         =     jit(model.generate      )
														UpperCAmelCase         =     jit_generate(UpperCamelCase__      ).sequences

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






		def SCREAMING_SNAKE_CASE_			(						self      :     str      )				->					int:
								'''simple docstring'''







								UpperCAmelCase       ,		UpperCAmelCase       ,		UpperCAmelCase       ,		UpperCAmelCase         =     self._get_input_ids_and_config()
								UpperCAmelCase         =     True
								UpperCAmelCase         =     max_length

								for model_class in self.all_generative_model_classes:
														UpperCAmelCase         =     model_class(UpperCamelCase__      )

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

														UpperCAmelCase         =     jit(model.generate      )
														UpperCAmelCase         =     jit_generate(UpperCamelCase__      ).sequences

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






		def SCREAMING_SNAKE_CASE_			(						self      :     int      )				->					List[Any]:
								'''simple docstring'''







								UpperCAmelCase       ,		UpperCAmelCase       ,		UpperCAmelCase       ,		UpperCAmelCase         =     self._get_input_ids_and_config()
								UpperCAmelCase         =     False
								UpperCAmelCase         =     max_length
								UpperCAmelCase         =     2

								for model_class in self.all_generative_model_classes:
														UpperCAmelCase         =     model_class(UpperCamelCase__      )

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

														UpperCAmelCase         =     jit(model.generate      )
														UpperCAmelCase         =     jit_generate(UpperCamelCase__      ).sequences

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






		def SCREAMING_SNAKE_CASE_			(						self      :     Tuple      )				->					Dict:
								'''simple docstring'''







								UpperCAmelCase       ,		UpperCAmelCase       ,		UpperCAmelCase       ,		UpperCAmelCase         =     self._get_input_ids_and_config()
								UpperCAmelCase         =     False
								UpperCAmelCase         =     max_length
								UpperCAmelCase         =     2
								UpperCAmelCase         =     2

								for model_class in self.all_generative_model_classes:
														UpperCAmelCase         =     model_class(UpperCamelCase__      )

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






		def SCREAMING_SNAKE_CASE_			(						self      :     Optional[int]      )				->					Any:
								'''simple docstring'''







								UpperCAmelCase       ,		UpperCAmelCase       ,		UpperCAmelCase       ,		UpperCAmelCase         =     self._get_input_ids_and_config()
								UpperCAmelCase         =     True
								UpperCAmelCase         =     max_length
								UpperCAmelCase         =     0.8
								UpperCAmelCase         =     10
								UpperCAmelCase         =     0.3
								UpperCAmelCase         =     1
								UpperCAmelCase         =     8
								UpperCAmelCase         =     9

								for model_class in self.all_generative_model_classes:
														UpperCAmelCase         =     model_class(UpperCamelCase__      )

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

														UpperCAmelCase         =     jit(model.generate      )
														UpperCAmelCase         =     jit_generate(UpperCamelCase__      ).sequences

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






		def SCREAMING_SNAKE_CASE_			(						self      :     Optional[int]      )				->					Union[str, Any]:
								'''simple docstring'''







								UpperCAmelCase       ,		UpperCAmelCase       ,		UpperCAmelCase       ,		UpperCAmelCase         =     self._get_input_ids_and_config()
								UpperCAmelCase         =     max_length
								UpperCAmelCase         =     1
								UpperCAmelCase         =     8
								UpperCAmelCase         =     9

								for model_class in self.all_generative_model_classes:
														UpperCAmelCase         =     model_class(UpperCamelCase__      )

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

														UpperCAmelCase         =     jit(model.generate      )
														UpperCAmelCase         =     jit_generate(UpperCamelCase__      ).sequences

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






		def SCREAMING_SNAKE_CASE_			(						self      :     int      )				->					str:
								'''simple docstring'''







								UpperCAmelCase       ,		UpperCAmelCase       ,		UpperCAmelCase       ,		UpperCAmelCase         =     self._get_input_ids_and_config()
								UpperCAmelCase         =     max_length
								UpperCAmelCase         =     2
								UpperCAmelCase         =     1
								UpperCAmelCase         =     8
								UpperCAmelCase         =     9

								for model_class in self.all_generative_model_classes:
														UpperCAmelCase         =     model_class(UpperCamelCase__      )

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

														UpperCAmelCase         =     jit(model.generate      )
														UpperCAmelCase         =     jit_generate(UpperCamelCase__      ).sequences

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






		def SCREAMING_SNAKE_CASE_			(						self      :     Dict      )				->					int:
								'''simple docstring'''







								UpperCAmelCase       ,		UpperCAmelCase       ,		UpperCAmelCase       ,		UpperCAmelCase         =     self._get_input_ids_and_config()

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

								UpperCAmelCase         =     False
								UpperCAmelCase         =     max_length

								for model_class in self.all_generative_model_classes:
														UpperCAmelCase         =     model_class(UpperCamelCase__      )

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

														UpperCAmelCase         =     jit(model.generate      )
														UpperCAmelCase         =     jit_generate(UpperCamelCase__				,		attention_mask=UpperCamelCase__      ).sequences

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






		def SCREAMING_SNAKE_CASE_			(						self      :     Dict      )				->					List[str]:
								'''simple docstring'''







								UpperCAmelCase       ,		UpperCAmelCase       ,		UpperCAmelCase       ,		UpperCAmelCase         =     self._get_input_ids_and_config()

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

								UpperCAmelCase         =     True
								UpperCAmelCase         =     max_length

								for model_class in self.all_generative_model_classes:
														UpperCAmelCase         =     model_class(UpperCamelCase__      )

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

														UpperCAmelCase         =     jit(model.generate      )
														UpperCAmelCase         =     jit_generate(UpperCamelCase__				,		attention_mask=UpperCamelCase__      ).sequences

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







		def SCREAMING_SNAKE_CASE_			(						self      :     str      )				->					List[str]:
								'''simple docstring'''







								UpperCAmelCase       ,		UpperCAmelCase       ,		UpperCAmelCase       ,		UpperCAmelCase         =     self._get_input_ids_and_config()

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

								UpperCAmelCase         =     2
								UpperCAmelCase         =     max_length

								for model_class in self.all_generative_model_classes:
														UpperCAmelCase         =     model_class(UpperCamelCase__      )

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

														UpperCAmelCase         =     jit(model.generate      )
														UpperCAmelCase         =     jit_generate(UpperCamelCase__				,		attention_mask=UpperCamelCase__      ).sequences

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





@require_flax
class  __magic_name__   (							unittest.TestCase						):






		def SCREAMING_SNAKE_CASE_			(						self      :     Dict      )				->					Optional[Any]:
								'''simple docstring'''







								UpperCAmelCase         =     AutoTokenizer.from_pretrained("hf-internal-testing/tiny-bert"      )
								UpperCAmelCase         =     FlaxAutoModelForCausalLM.from_pretrained("hf-internal-testing/tiny-bert-flax-only"      )

								UpperCAmelCase         =     "Hello world"
								UpperCAmelCase         =     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"      ):
														UpperCAmelCase         =     {"foo": "bar"}
														model.generate(UpperCamelCase__				,		**UpperCamelCase__      )








import unittest

import torch
from torch import nn

from diffusers.models.activations import get_activation





class  __magic_name__   (							unittest.TestCase						):






		def SCREAMING_SNAKE_CASE_			(						self      :     int      )				->					str:
								'''simple docstring'''







								UpperCAmelCase         =     get_activation("swish"      )

								self.assertIsInstance(UpperCamelCase__				,		nn.SiLU      )

								self.assertEqual(act(torch.tensor(-1_00				,		dtype=torch.floataa      )      ).item()				,		0      )
								self.assertNotEqual(act(torch.tensor(-1				,		dtype=torch.floataa      )      ).item()				,		0      )
								self.assertEqual(act(torch.tensor(0				,		dtype=torch.floataa      )      ).item()				,		0      )
								self.assertEqual(act(torch.tensor(20				,		dtype=torch.floataa      )      ).item()				,		20      )






		def SCREAMING_SNAKE_CASE_			(						self      :     Tuple      )				->					List[Any]:
								'''simple docstring'''







								UpperCAmelCase         =     get_activation("silu"      )

								self.assertIsInstance(UpperCamelCase__				,		nn.SiLU      )

								self.assertEqual(act(torch.tensor(-1_00				,		dtype=torch.floataa      )      ).item()				,		0      )
								self.assertNotEqual(act(torch.tensor(-1				,		dtype=torch.floataa      )      ).item()				,		0      )
								self.assertEqual(act(torch.tensor(0				,		dtype=torch.floataa      )      ).item()				,		0      )
								self.assertEqual(act(torch.tensor(20				,		dtype=torch.floataa      )      ).item()				,		20      )






		def SCREAMING_SNAKE_CASE_			(						self      :     Optional[Any]      )				->					int:
								'''simple docstring'''







								UpperCAmelCase         =     get_activation("mish"      )

								self.assertIsInstance(UpperCamelCase__				,		nn.Mish      )

								self.assertEqual(act(torch.tensor(-2_00				,		dtype=torch.floataa      )      ).item()				,		0      )
								self.assertNotEqual(act(torch.tensor(-1				,		dtype=torch.floataa      )      ).item()				,		0      )
								self.assertEqual(act(torch.tensor(0				,		dtype=torch.floataa      )      ).item()				,		0      )
								self.assertEqual(act(torch.tensor(20				,		dtype=torch.floataa      )      ).item()				,		20      )







		def SCREAMING_SNAKE_CASE_			(						self      :     str      )				->					List[str]:
								'''simple docstring'''







								UpperCAmelCase         =     get_activation("gelu"      )

								self.assertIsInstance(UpperCamelCase__				,		nn.GELU      )

								self.assertEqual(act(torch.tensor(-1_00				,		dtype=torch.floataa      )      ).item()				,		0      )
								self.assertNotEqual(act(torch.tensor(-1				,		dtype=torch.floataa      )      ).item()				,		0      )
								self.assertEqual(act(torch.tensor(0				,		dtype=torch.floataa      )      ).item()				,		0      )
								self.assertEqual(act(torch.tensor(20				,		dtype=torch.floataa      )      ).item()				,		20      )



