




import gc
import unittest

from parameterized import parameterized

from diffusers import FlaxUNetaDConditionModel
from diffusers.utils import is_flax_available
from diffusers.utils.testing_utils import load_hf_numpy, require_flax, slow


if is_flax_available():
		import jax
		import jax.numpy as jnp
@slow
@require_flax
class 		UpperCAmelCase   (				unittest.TestCase			):







						'''simple docstring'''





						def        snake_case__						(						self							:     Dict       ,	__lowercase							:     int       ,	__lowercase							:     Tuple						):






							"""simple docstring"""






							return f"gaussian_noise_s={seed}_shape={'_'.join([str(__lowercase						) for s in shape]						)}.npy"






						def        snake_case__						(						self							:     List[Any]						):






							"""simple docstring"""






							super().tearDown()
							gc.collect()






						def        snake_case__						(						self							:     int       ,	__lowercase							:     List[str]=0       ,	__lowercase							:     Dict=(4, 4, 64, 64)       ,	__lowercase							:     Union[str, Any]=False						):






							"""simple docstring"""






							snake_case_		    =     jnp.bfloataa if fpaa else jnp.floataa
							snake_case_		    =     jnp.array(load_hf_numpy(self.get_file_format(__lowercase       ,	__lowercase						)						)       ,	dtype=__lowercase						)
							return image






						def        snake_case__						(						self							:     Optional[Any]       ,	__lowercase							:     List[Any]=False       ,	__lowercase							:     Any="CompVis/stable-diffusion-v1-4"						):






							"""simple docstring"""






							snake_case_		    =     jnp.bfloataa if fpaa else jnp.floataa
							snake_case_		    =     "bf16" if fpaa else None

							snake_case_    ,      snake_case_		    =     FlaxUNetaDConditionModel.from_pretrained(
							    __lowercase       ,	subfolder="unet"       ,	dtype=__lowercase       ,	revision=__lowercase						)
							return model, params






						def        snake_case__						(						self							:     Any       ,	__lowercase							:     Dict=0       ,	__lowercase							:     Optional[Any]=(4, 77, 7_68)       ,	__lowercase							:     List[str]=False						):






							"""simple docstring"""






							snake_case_		    =     jnp.bfloataa if fpaa else jnp.floataa
							snake_case_		    =     jnp.array(load_hf_numpy(self.get_file_format(__lowercase       ,	__lowercase						)						)       ,	dtype=__lowercase						)
							return hidden_states






						@parameterized.expand(
						    [
						        # fmt: off
						        [83, 4, [-0.2323, -0.1304, 0.0813, -0.3093, -0.0919, -0.1571, -0.1125, -0.5806]],
						        [17, 0.55, [-0.0831, -0.2443, 0.0901, -0.0919, 0.3396, 0.0103, -0.3743, 0.0701]],
						        [8, 0.89, [-0.4863, 0.0859, 0.0875, -0.1658, 0.9199, -0.0114, 0.4839, 0.4639]],
						        [3, 10_00, [-0.5649, 0.2402, -0.5518, 0.1248, 1.1328, -0.2443, -0.0325, -1.0078]],
						        # fmt: on
						    ]						)
						def        snake_case__						(						self							:     List[str]       ,	__lowercase							:     Union[str, Any]       ,	__lowercase							:     Optional[int]       ,	__lowercase							:     str						):






							"""simple docstring"""






							snake_case_    ,      snake_case_		    =     self.get_unet_model(model_id="CompVis/stable-diffusion-v1-4"       ,	fpaa=__lowercase						)
							snake_case_		    =     self.get_latents(__lowercase       ,	fpaa=__lowercase						)
							snake_case_		    =     self.get_encoder_hidden_states(__lowercase       ,	fpaa=__lowercase						)

							snake_case_		    =     model.apply(
							    {"params": params}       ,	__lowercase       ,	jnp.array(__lowercase       ,	dtype=jnp.intaa						)       ,	encoder_hidden_states=__lowercase       ,	).sample

							assert sample.shape == latents.shape

							snake_case_		    =     jnp.asarray(jax.device_get((sample[-1, -2:, -2:, :2].flatten())						)       ,	dtype=jnp.floataa						)
							snake_case_		    =     jnp.array(__lowercase       ,	dtype=jnp.floataa						)

							# Found torch (float16) and flax (bfloat16) outputs to be within this tolerance, in the same hardware
							assert jnp.allclose(__lowercase       ,	__lowercase       ,	atol=1E-2						)






						@parameterized.expand(
						    [
						        # fmt: off
						        [83, 4, [0.1514, 0.0807, 0.1624, 0.1016, -0.1896, 0.0263, 0.0677, 0.2310]],
						        [17, 0.55, [0.1164, -0.0216, 0.0170, 0.1589, -0.3120, 0.1005, -0.0581, -0.1458]],
						        [8, 0.89, [-0.1758, -0.0169, 0.1004, -0.1411, 0.1312, 0.1103, -0.1996, 0.2139]],
						        [3, 10_00, [0.1214, 0.0352, -0.0731, -0.1562, -0.0994, -0.0906, -0.2340, -0.0539]],
						        # fmt: on
						    ]						)
						def        snake_case__						(						self							:     Optional[Any]       ,	__lowercase							:     Optional[int]       ,	__lowercase							:     List[str]       ,	__lowercase							:     Optional[int]						):






							"""simple docstring"""






							snake_case_    ,      snake_case_		    =     self.get_unet_model(model_id="stabilityai/stable-diffusion-2"       ,	fpaa=__lowercase						)
							snake_case_		    =     self.get_latents(__lowercase       ,	shape=(4, 4, 96, 96)       ,	fpaa=__lowercase						)
							snake_case_		    =     self.get_encoder_hidden_states(__lowercase       ,	shape=(4, 77, 10_24)       ,	fpaa=__lowercase						)

							snake_case_		    =     model.apply(
							    {"params": params}       ,	__lowercase       ,	jnp.array(__lowercase       ,	dtype=jnp.intaa						)       ,	encoder_hidden_states=__lowercase       ,	).sample

							assert sample.shape == latents.shape

							snake_case_		    =     jnp.asarray(jax.device_get((sample[-1, -2:, -2:, :2].flatten())						)       ,	dtype=jnp.floataa						)
							snake_case_		    =     jnp.array(__lowercase       ,	dtype=jnp.floataa						)

							# Found torch (float16) and flax (bfloat16) outputs to be within this tolerance, on the same hardware
							assert jnp.allclose(__lowercase       ,	__lowercase       ,	atol=1E-2						)






from timeit import timeit

lowercase__				:     Union[str, Any] 		=						{
    "MALAYALAM": True,
    "String": False,
    "rotor": True,
    "level": True,
    "A": True,
    "BB": True,
    "ABC": False,
    "amanaplanacanalpanama": True,  # "a man a plan a canal panama"
}
# Ensure our test data is valid
assert all((key == key[::-1]) is value for key, value in test_data.items())


def 						lowerCamelCase__    (		_A							):







	'''simple docstring'''





	snake_case_		    =     0
	snake_case_		    =     len(_A							) - 1
	while start_i < end_i:
		if s[start_i] == s[end_i]:
			start_i += 1
			end_i -= 1
		else:
			return False
	return True


def 						lowerCamelCase__    (		_A							):







	'''simple docstring'''





	snake_case_		    =     len(_A							) // 2
	snake_case_		    =     len(_A							)

	# We need to traverse till half of the length of string
	# as we can get access of the i'th last element from
	# i'th index.
	# eg: [0,1,2,3,4,5] => 4th index can be accessed
	# with the help of 1st index (i==n-i-1)
	# where n is length of string
	return all(s[i] == s[n - i - 1] for i in range(_A							)							)


def 						lowerCamelCase__    (		_A							):







	'''simple docstring'''





	if len(_A							) <= 2:
		return True
	if s[0] == s[len(_A							) - 1]:
		return is_palindrome_recursive(s[1:-1]							)
	else:
		return False


def 						lowerCamelCase__    (		_A							):







	'''simple docstring'''





	return s == s[::-1]


def 						lowerCamelCase__    (		_A							):







	'''simple docstring'''





	snake_case_		    =     f"all({name}(key) is value for key, value in test_data.items())"
	snake_case_		    =     f"from __main__ import test_data, {name}"
	snake_case_		    =     500000
	snake_case_		    =     timeit(stmt=_A     ,	setup=_A     ,	number=_A							)
	print(f"{name:<35} finished {number:,} runs in {result:.5f} seconds"							)


if __name__ == "__main__":
		for key, value in test_data.items():
				assert is_palindrome(key) is is_palindrome_recursive(key)
				assert is_palindrome(key) is is_palindrome_slice(key)
				print(f'''{key:21} {value}''')
		print("a man a plan a canal panama")

		# finished 500,000 runs in 0.46793 seconds
		benchmark_function("is_palindrome_slice")
		# finished 500,000 runs in 0.85234 seconds
		benchmark_function("is_palindrome")
		# finished 500,000 runs in 1.32028 seconds
		benchmark_function("is_palindrome_recursive")
		# finished 500,000 runs in 2.08679 seconds
		benchmark_function("is_palindrome_traversal")
