


'''simple docstring'''



from collections import deque
from math import floor
from random import random
from time import time





class    UpperCAmelCase_  :





    """simple docstring"""






    def __init__(							self	)				->     List[str]:
         __lowerCamelCase   :	Optional[int]								=		{}







    def       lowercase_				(							self			,     SCREAMING_SNAKE_CASE_			,     SCREAMING_SNAKE_CASE_			,     SCREAMING_SNAKE_CASE_=1	)				->     str:
         if self.graph.get(SCREAMING_SNAKE_CASE_	):
              if self.graph[u].count([w, v]	) == 0:
                   self.graph[u].append([w, v]	)
         else:
              __lowerCamelCase   :	List[Any]								=		[[w, v]]
         if not self.graph.get(SCREAMING_SNAKE_CASE_	):
              __lowerCamelCase   :	Optional[int]								=		[]







    def       lowercase_				(							self	)				->     Union[str, Any]:
         return list(self.graph	)







    def       lowercase_				(							self			,     SCREAMING_SNAKE_CASE_			,     SCREAMING_SNAKE_CASE_	)				->     Optional[Any]:
         if self.graph.get(SCREAMING_SNAKE_CASE_	):
              for _ in self.graph[u]:
                   if _[1] == v:
                        self.graph[u].remove(SCREAMING_SNAKE_CASE_	)







    def       lowercase_				(							self			,     SCREAMING_SNAKE_CASE_=-2			,     SCREAMING_SNAKE_CASE_=-1	)				->     Dict:
         if s == d:
              return []
         __lowerCamelCase   :	Dict								=		[]
         __lowerCamelCase   :	List[str]								=		[]
         if s == -2:
              __lowerCamelCase   :	Any								=		list(self.graph	)[0]
         stack.append(SCREAMING_SNAKE_CASE_	)
         visited.append(SCREAMING_SNAKE_CASE_	)
         __lowerCamelCase   :	Optional[Any]								=		s

         while True:
              # check if there is any non isolated nodes
              if len(self.graph[s]	) != 0:
                   __lowerCamelCase   :	str								=		s
                   for node in self.graph[s]:
                        if visited.count(node[1]	) < 1:
                             if node[1] == d:
                                  visited.append(SCREAMING_SNAKE_CASE_	)
                                  return visited
                             else:
                                  stack.append(node[1]	)
                                  visited.append(node[1]	)
                                  __lowerCamelCase   :	Optional[int]								=		node[1]
                                  break

            # check if all the children are visited
              if s == ss:
                   stack.pop()
                   if len(SCREAMING_SNAKE_CASE_	) != 0:
                        __lowerCamelCase   :	List[Any]								=		stack[len(SCREAMING_SNAKE_CASE_	) - 1]
              else:
                   __lowerCamelCase   :	List[Any]								=		ss

              # check if se have reached the starting point
              if len(SCREAMING_SNAKE_CASE_	) == 0:
                   return visited







    def       lowercase_				(							self			,     SCREAMING_SNAKE_CASE_=-1	)				->     str:
         if c == -1:
              __lowerCamelCase   :	Dict								=		floor(random() * 1_00_00	) + 10
         for i in range(SCREAMING_SNAKE_CASE_	):
              # every vertex has max 100 edges
              for _ in range(floor(random() * 1_02	) + 1	):
                   __lowerCamelCase   :	Tuple								=		floor(random() * c	) + 1
                   if n != i:
                        self.add_pair(SCREAMING_SNAKE_CASE_			,     SCREAMING_SNAKE_CASE_			,     1	)







    def       lowercase_				(							self			,     SCREAMING_SNAKE_CASE_=-2	)				->     Dict:
         __lowerCamelCase   :	Union[str, Any]								=		deque()
         __lowerCamelCase   :	Any								=		[]
         if s == -2:
              __lowerCamelCase   :	str								=		list(self.graph	)[0]
         d.append(SCREAMING_SNAKE_CASE_	)
         visited.append(SCREAMING_SNAKE_CASE_	)
         while d:
              __lowerCamelCase   :	Any								=		d.popleft()
              if len(self.graph[s]	) != 0:
                   for node in self.graph[s]:
                        if visited.count(node[1]	) < 1:
                             d.append(node[1]	)
                             visited.append(node[1]	)
         return visited







    def       lowercase_				(							self			,     SCREAMING_SNAKE_CASE_	)				->     Optional[int]:
         __lowerCamelCase   :	Union[str, Any]								=		0
         for x in self.graph:
              for y in self.graph[x]:
                   if y[1] == u:
                        count += 1
         return count







    def       lowercase_				(							self			,     SCREAMING_SNAKE_CASE_	)				->     Optional[int]:
         return len(self.graph[u]	)







    def       lowercase_				(							self			,     SCREAMING_SNAKE_CASE_=-2	)				->     Any:
         __lowerCamelCase   :	str								=		[]
         __lowerCamelCase   :	Any								=		[]
         if s == -2:
              __lowerCamelCase   :	Any								=		list(self.graph	)[0]
         stack.append(SCREAMING_SNAKE_CASE_	)
         visited.append(SCREAMING_SNAKE_CASE_	)
         __lowerCamelCase   :	Union[str, Any]								=		s
         __lowerCamelCase   :	Tuple								=		[]

         while True:
              # check if there is any non isolated nodes
              if len(self.graph[s]	) != 0:
                   __lowerCamelCase   :	Dict								=		s
                   for node in self.graph[s]:
                        if visited.count(node[1]	) < 1:
                             stack.append(node[1]	)
                             visited.append(node[1]	)
                             __lowerCamelCase   :	Tuple								=		node[1]
                             break

            # check if all the children are visited
              if s == ss:
                   sorted_nodes.append(stack.pop()	)
                   if len(SCREAMING_SNAKE_CASE_	) != 0:
                        __lowerCamelCase   :	str								=		stack[len(SCREAMING_SNAKE_CASE_	) - 1]
              else:
                   __lowerCamelCase   :	List[Any]								=		ss

              # check if se have reached the starting point
              if len(SCREAMING_SNAKE_CASE_	) == 0:
                   return sorted_nodes







    def       lowercase_				(							self	)				->     List[str]:
         __lowerCamelCase   :	str								=		[]
         __lowerCamelCase   :	Dict								=		[]
         __lowerCamelCase   :	int								=		list(self.graph	)[0]
         stack.append(SCREAMING_SNAKE_CASE_	)
         visited.append(SCREAMING_SNAKE_CASE_	)
         __lowerCamelCase   :	List[Any]								=		-2
         __lowerCamelCase   :	Any								=		[]
         __lowerCamelCase   :	List[Any]								=		s
         __lowerCamelCase   :	List[Any]								=		False
         __lowerCamelCase   :	List[str]								=		set()

         while True:
              # check if there is any non isolated nodes
              if len(self.graph[s]	) != 0:
                   __lowerCamelCase   :	str								=		s
                   for node in self.graph[s]:
                        if (
                            visited.count(node[1]	) > 0
                            and node[1] != parent
                            and indirect_parents.count(node[1]	) > 0
                            and not on_the_way_back
                        ):
                             __lowerCamelCase   :	Union[str, Any]								=		len(SCREAMING_SNAKE_CASE_	) - 1
                             while len_stack >= 0:
                                  if stack[len_stack] == node[1]:
                                       anticipating_nodes.add(node[1]	)
                                       break
                                  else:
                                       anticipating_nodes.add(stack[len_stack]	)
                                       len_stack -= 1
                        if visited.count(node[1]	) < 1:
                             stack.append(node[1]	)
                             visited.append(node[1]	)
                             __lowerCamelCase   :	List[Any]								=		node[1]
                             break

            # check if all the children are visited
              if s == ss:
                   stack.pop()
                   __lowerCamelCase   :	List[str]								=		True
                   if len(SCREAMING_SNAKE_CASE_	) != 0:
                        __lowerCamelCase   :	Dict								=		stack[len(SCREAMING_SNAKE_CASE_	) - 1]
              else:
                   __lowerCamelCase   :	Any								=		False
                   indirect_parents.append(SCREAMING_SNAKE_CASE_	)
                   __lowerCamelCase   :	Any								=		s
                   __lowerCamelCase   :	Optional[Any]								=		ss

              # check if se have reached the starting point
              if len(SCREAMING_SNAKE_CASE_	) == 0:
                   return list(SCREAMING_SNAKE_CASE_	)







    def       lowercase_				(							self	)				->     Any:
         __lowerCamelCase   :	Union[str, Any]								=		[]
         __lowerCamelCase   :	List[str]								=		[]
         __lowerCamelCase   :	Optional[int]								=		list(self.graph	)[0]
         stack.append(SCREAMING_SNAKE_CASE_	)
         visited.append(SCREAMING_SNAKE_CASE_	)
         __lowerCamelCase   :	Dict								=		-2
         __lowerCamelCase   :	Optional[int]								=		[]
         __lowerCamelCase   :	Any								=		s
         __lowerCamelCase   :	Dict								=		False
         __lowerCamelCase   :	Any								=		set()

         while True:
              # check if there is any non isolated nodes
              if len(self.graph[s]	) != 0:
                   __lowerCamelCase   :	List[Any]								=		s
                   for node in self.graph[s]:
                        if (
                            visited.count(node[1]	) > 0
                            and node[1] != parent
                            and indirect_parents.count(node[1]	) > 0
                            and not on_the_way_back
                        ):
                             __lowerCamelCase   :	str								=		len(SCREAMING_SNAKE_CASE_	) - 1
                             while len_stack_minus_one >= 0:
                                  if stack[len_stack_minus_one] == node[1]:
                                       anticipating_nodes.add(node[1]	)
                                       break
                                  else:
                                       return True
                        if visited.count(node[1]	) < 1:
                             stack.append(node[1]	)
                             visited.append(node[1]	)
                             __lowerCamelCase   :	Union[str, Any]								=		node[1]
                             break

            # check if all the children are visited
              if s == ss:
                   stack.pop()
                   __lowerCamelCase   :	Tuple								=		True
                   if len(SCREAMING_SNAKE_CASE_	) != 0:
                        __lowerCamelCase   :	Optional[int]								=		stack[len(SCREAMING_SNAKE_CASE_	) - 1]
              else:
                   __lowerCamelCase   :	str								=		False
                   indirect_parents.append(SCREAMING_SNAKE_CASE_	)
                   __lowerCamelCase   :	Any								=		s
                   __lowerCamelCase   :	str								=		ss

              # check if se have reached the starting point
              if len(SCREAMING_SNAKE_CASE_	) == 0:
                   return False







    def       lowercase_				(							self			,     SCREAMING_SNAKE_CASE_=-2			,     SCREAMING_SNAKE_CASE_=-1	)				->     Optional[Any]:
         __lowerCamelCase   :	Optional[Any]								=		time()
         self.dfs(SCREAMING_SNAKE_CASE_			,     SCREAMING_SNAKE_CASE_	)
         __lowerCamelCase   :	int								=		time()
         return end - begin







    def       lowercase_				(							self			,     SCREAMING_SNAKE_CASE_=-2	)				->     Optional[Any]:
         __lowerCamelCase   :	Optional[int]								=		time()
         self.bfs(SCREAMING_SNAKE_CASE_	)
         __lowerCamelCase   :	str								=		time()
         return end - begin







class    UpperCAmelCase_  :





    """simple docstring"""






    def __init__(							self	)				->     Optional[int]:
         __lowerCamelCase   :	List[Any]								=		{}







    def       lowercase_				(							self			,     SCREAMING_SNAKE_CASE_			,     SCREAMING_SNAKE_CASE_			,     SCREAMING_SNAKE_CASE_=1	)				->     Union[str, Any]:
         # check if the u exists
         if self.graph.get(SCREAMING_SNAKE_CASE_	):
              # if there already is a edge
              if self.graph[u].count([w, v]	) == 0:
                   self.graph[u].append([w, v]	)
         else:
              # if u does not exist
              __lowerCamelCase   :	int								=		[[w, v]]
         # add the other way
         if self.graph.get(SCREAMING_SNAKE_CASE_	):
              # if there already is a edge
              if self.graph[v].count([w, u]	) == 0:
                   self.graph[v].append([w, u]	)
         else:
              # if u does not exist
              __lowerCamelCase   :	Optional[int]								=		[[w, u]]







    def       lowercase_				(							self			,     SCREAMING_SNAKE_CASE_			,     SCREAMING_SNAKE_CASE_	)				->     str:
         if self.graph.get(SCREAMING_SNAKE_CASE_	):
              for _ in self.graph[u]:
                   if _[1] == v:
                        self.graph[u].remove(SCREAMING_SNAKE_CASE_	)
        # the other way round
         if self.graph.get(SCREAMING_SNAKE_CASE_	):
              for _ in self.graph[v]:
                   if _[1] == u:
                        self.graph[v].remove(SCREAMING_SNAKE_CASE_	)







    def       lowercase_				(							self			,     SCREAMING_SNAKE_CASE_=-2			,     SCREAMING_SNAKE_CASE_=-1	)				->     Union[str, Any]:
         if s == d:
              return []
         __lowerCamelCase   :	Optional[Any]								=		[]
         __lowerCamelCase   :	Tuple								=		[]
         if s == -2:
              __lowerCamelCase   :	Any								=		list(self.graph	)[0]
         stack.append(SCREAMING_SNAKE_CASE_	)
         visited.append(SCREAMING_SNAKE_CASE_	)
         __lowerCamelCase   :	str								=		s

         while True:
              # check if there is any non isolated nodes
              if len(self.graph[s]	) != 0:
                   __lowerCamelCase   :	Dict								=		s
                   for node in self.graph[s]:
                        if visited.count(node[1]	) < 1:
                             if node[1] == d:
                                  visited.append(SCREAMING_SNAKE_CASE_	)
                                  return visited
                             else:
                                  stack.append(node[1]	)
                                  visited.append(node[1]	)
                                  __lowerCamelCase   :	Union[str, Any]								=		node[1]
                                  break

            # check if all the children are visited
              if s == ss:
                   stack.pop()
                   if len(SCREAMING_SNAKE_CASE_	) != 0:
                        __lowerCamelCase   :	Optional[int]								=		stack[len(SCREAMING_SNAKE_CASE_	) - 1]
              else:
                   __lowerCamelCase   :	Tuple								=		ss

              # check if se have reached the starting point
              if len(SCREAMING_SNAKE_CASE_	) == 0:
                   return visited







    def       lowercase_				(							self			,     SCREAMING_SNAKE_CASE_=-1	)				->     str:
         if c == -1:
              __lowerCamelCase   :	Union[str, Any]								=		floor(random() * 1_00_00	) + 10
         for i in range(SCREAMING_SNAKE_CASE_	):
              # every vertex has max 100 edges
              for _ in range(floor(random() * 1_02	) + 1	):
                   __lowerCamelCase   :	str								=		floor(random() * c	) + 1
                   if n != i:
                        self.add_pair(SCREAMING_SNAKE_CASE_			,     SCREAMING_SNAKE_CASE_			,     1	)







    def       lowercase_				(							self			,     SCREAMING_SNAKE_CASE_=-2	)				->     List[str]:
         __lowerCamelCase   :	Tuple								=		deque()
         __lowerCamelCase   :	List[Any]								=		[]
         if s == -2:
              __lowerCamelCase   :	int								=		list(self.graph	)[0]
         d.append(SCREAMING_SNAKE_CASE_	)
         visited.append(SCREAMING_SNAKE_CASE_	)
         while d:
              __lowerCamelCase   :	Tuple								=		d.popleft()
              if len(self.graph[s]	) != 0:
                   for node in self.graph[s]:
                        if visited.count(node[1]	) < 1:
                             d.append(node[1]	)
                             visited.append(node[1]	)
         return visited







    def       lowercase_				(							self			,     SCREAMING_SNAKE_CASE_	)				->     Optional[Any]:
         return len(self.graph[u]	)







    def       lowercase_				(							self	)				->     int:
         __lowerCamelCase   :	Union[str, Any]								=		[]
         __lowerCamelCase   :	Optional[Any]								=		[]
         __lowerCamelCase   :	Any								=		list(self.graph	)[0]
         stack.append(SCREAMING_SNAKE_CASE_	)
         visited.append(SCREAMING_SNAKE_CASE_	)
         __lowerCamelCase   :	Tuple								=		-2
         __lowerCamelCase   :	Optional[Any]								=		[]
         __lowerCamelCase   :	int								=		s
         __lowerCamelCase   :	Tuple								=		False
         __lowerCamelCase   :	Dict								=		set()

         while True:
              # check if there is any non isolated nodes
              if len(self.graph[s]	) != 0:
                   __lowerCamelCase   :	Any								=		s
                   for node in self.graph[s]:
                        if (
                            visited.count(node[1]	) > 0
                            and node[1] != parent
                            and indirect_parents.count(node[1]	) > 0
                            and not on_the_way_back
                        ):
                             __lowerCamelCase   :	Union[str, Any]								=		len(SCREAMING_SNAKE_CASE_	) - 1
                             while len_stack >= 0:
                                  if stack[len_stack] == node[1]:
                                       anticipating_nodes.add(node[1]	)
                                       break
                                  else:
                                       anticipating_nodes.add(stack[len_stack]	)
                                       len_stack -= 1
                        if visited.count(node[1]	) < 1:
                             stack.append(node[1]	)
                             visited.append(node[1]	)
                             __lowerCamelCase   :	Tuple								=		node[1]
                             break

            # check if all the children are visited
              if s == ss:
                   stack.pop()
                   __lowerCamelCase   :	Dict								=		True
                   if len(SCREAMING_SNAKE_CASE_	) != 0:
                        __lowerCamelCase   :	Any								=		stack[len(SCREAMING_SNAKE_CASE_	) - 1]
              else:
                   __lowerCamelCase   :	str								=		False
                   indirect_parents.append(SCREAMING_SNAKE_CASE_	)
                   __lowerCamelCase   :	int								=		s
                   __lowerCamelCase   :	Tuple								=		ss

              # check if se have reached the starting point
              if len(SCREAMING_SNAKE_CASE_	) == 0:
                   return list(SCREAMING_SNAKE_CASE_	)







    def       lowercase_				(							self	)				->     int:
         __lowerCamelCase   :	Dict								=		[]
         __lowerCamelCase   :	Union[str, Any]								=		[]
         __lowerCamelCase   :	Union[str, Any]								=		list(self.graph	)[0]
         stack.append(SCREAMING_SNAKE_CASE_	)
         visited.append(SCREAMING_SNAKE_CASE_	)
         __lowerCamelCase   :	Optional[int]								=		-2
         __lowerCamelCase   :	Dict								=		[]
         __lowerCamelCase   :	List[Any]								=		s
         __lowerCamelCase   :	Tuple								=		False
         __lowerCamelCase   :	Optional[Any]								=		set()

         while True:
              # check if there is any non isolated nodes
              if len(self.graph[s]	) != 0:
                   __lowerCamelCase   :	Optional[Any]								=		s
                   for node in self.graph[s]:
                        if (
                            visited.count(node[1]	) > 0
                            and node[1] != parent
                            and indirect_parents.count(node[1]	) > 0
                            and not on_the_way_back
                        ):
                             __lowerCamelCase   :	Optional[int]								=		len(SCREAMING_SNAKE_CASE_	) - 1
                             while len_stack_minus_one >= 0:
                                  if stack[len_stack_minus_one] == node[1]:
                                       anticipating_nodes.add(node[1]	)
                                       break
                                  else:
                                       return True
                        if visited.count(node[1]	) < 1:
                             stack.append(node[1]	)
                             visited.append(node[1]	)
                             __lowerCamelCase   :	Optional[int]								=		node[1]
                             break

            # check if all the children are visited
              if s == ss:
                   stack.pop()
                   __lowerCamelCase   :	Optional[int]								=		True
                   if len(SCREAMING_SNAKE_CASE_	) != 0:
                        __lowerCamelCase   :	List[str]								=		stack[len(SCREAMING_SNAKE_CASE_	) - 1]
              else:
                   __lowerCamelCase   :	Optional[Any]								=		False
                   indirect_parents.append(SCREAMING_SNAKE_CASE_	)
                   __lowerCamelCase   :	str								=		s
                   __lowerCamelCase   :	Tuple								=		ss

              # check if se have reached the starting point
              if len(SCREAMING_SNAKE_CASE_	) == 0:
                   return False







    def       lowercase_				(							self	)				->     Dict:
         return list(self.graph	)







    def       lowercase_				(							self			,     SCREAMING_SNAKE_CASE_=-2			,     SCREAMING_SNAKE_CASE_=-1	)				->     Union[str, Any]:
         __lowerCamelCase   :	Optional[int]								=		time()
         self.dfs(SCREAMING_SNAKE_CASE_			,     SCREAMING_SNAKE_CASE_	)
         __lowerCamelCase   :	List[str]								=		time()
         return end - begin







    def       lowercase_				(							self			,     SCREAMING_SNAKE_CASE_=-2	)				->     Optional[int]:
         __lowerCamelCase   :	Any								=		time()
         self.bfs(SCREAMING_SNAKE_CASE_	)
         __lowerCamelCase   :	int								=		time()
         return end - begin










'''simple docstring'''



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       lowercase_				(							self			,     SCREAMING_SNAKE_CASE_			,     SCREAMING_SNAKE_CASE_	)				->     Optional[int]:
         return f'gaussian_noise_s={seed}_shape={"_".join([str(SCREAMING_SNAKE_CASE_	) for s in shape]	)}.npy'







    def       lowercase_				(							self	)				->     Optional[Any]:
         # clean up the VRAM after each test
         super().tearDown()
         gc.collect()







    def       lowercase_				(							self			,     SCREAMING_SNAKE_CASE_=0			,     SCREAMING_SNAKE_CASE_=(4, 4, 64, 64)			,     SCREAMING_SNAKE_CASE_=False	)				->     str:
         __lowerCamelCase   :	List[str]								=		jnp.bfloataa if fpaa else jnp.floataa
         __lowerCamelCase   :	Tuple								=		jnp.array(load_hf_numpy(self.get_file_format(SCREAMING_SNAKE_CASE_			,     SCREAMING_SNAKE_CASE_	)	)			,     dtype=SCREAMING_SNAKE_CASE_	)
         return image







    def       lowercase_				(							self			,     SCREAMING_SNAKE_CASE_=False			,     SCREAMING_SNAKE_CASE_="CompVis/stable-diffusion-v1-4"	)				->     Dict:
         __lowerCamelCase   :	Union[str, Any]								=		jnp.bfloataa if fpaa else jnp.floataa
         __lowerCamelCase   :	Optional[Any]								=		'bf16' if fpaa else None

         __lowerCamelCase				,       __lowerCamelCase   :	str								=		FlaxUNetaDConditionModel.from_pretrained(
             SCREAMING_SNAKE_CASE_			,     subfolder='unet'			,     dtype=SCREAMING_SNAKE_CASE_			,     revision=SCREAMING_SNAKE_CASE_	)
         return model, params







    def       lowercase_				(							self			,     SCREAMING_SNAKE_CASE_=0			,     SCREAMING_SNAKE_CASE_=(4, 77, 7_68)			,     SCREAMING_SNAKE_CASE_=False	)				->     Union[str, Any]:
         __lowerCamelCase   :	Any								=		jnp.bfloataa if fpaa else jnp.floataa
         __lowerCamelCase   :	Optional[Any]								=		jnp.array(load_hf_numpy(self.get_file_format(SCREAMING_SNAKE_CASE_			,     SCREAMING_SNAKE_CASE_	)	)			,     dtype=SCREAMING_SNAKE_CASE_	)
         return hidden_states







    @parameterized.expand(
        [
            # fmt: off
            [83, 4, [-0.2_3_2_3, -0.1_3_0_4, 0.0_8_1_3, -0.3_0_9_3, -0.0_9_1_9, -0.1_5_7_1, -0.1_1_2_5, -0.5_8_0_6]],
            [17, 0.5_5, [-0.0_8_3_1, -0.2_4_4_3, 0.0_9_0_1, -0.0_9_1_9, 0.3_3_9_6, 0.0_1_0_3, -0.3_7_4_3, 0.0_7_0_1]],
            [8, 0.8_9, [-0.4_8_6_3, 0.0_8_5_9, 0.0_8_7_5, -0.1_6_5_8, 0.9_1_9_9, -0.0_1_1_4, 0.4_8_3_9, 0.4_6_3_9]],
            [3, 10_00, [-0.5_6_4_9, 0.2_4_0_2, -0.5_5_1_8, 0.1_2_4_8, 1.1_3_2_8, -0.2_4_4_3, -0.0_3_2_5, -1.0_0_7_8]],
            # fmt: on
        ]	)
    def       lowercase_				(							self			,     SCREAMING_SNAKE_CASE_			,     SCREAMING_SNAKE_CASE_			,     SCREAMING_SNAKE_CASE_	)				->     Any:
         __lowerCamelCase				,       __lowerCamelCase   :	Tuple								=		self.get_unet_model(model_id='CompVis/stable-diffusion-v1-4'			,     fpaa=SCREAMING_SNAKE_CASE_	)
         __lowerCamelCase   :	Any								=		self.get_latents(SCREAMING_SNAKE_CASE_			,     fpaa=SCREAMING_SNAKE_CASE_	)
         __lowerCamelCase   :	int								=		self.get_encoder_hidden_states(SCREAMING_SNAKE_CASE_			,     fpaa=SCREAMING_SNAKE_CASE_	)

         __lowerCamelCase   :	List[Any]								=		model.apply(
             {'params': params}			,     SCREAMING_SNAKE_CASE_			,     jnp.array(SCREAMING_SNAKE_CASE_			,     dtype=jnp.intaa	)			,     encoder_hidden_states=SCREAMING_SNAKE_CASE_			,     ).sample

         assert sample.shape == latents.shape

         __lowerCamelCase   :	Any								=		jnp.asarray(jax.device_get((sample[-1, -2:, -2:, :2].flatten())	)			,     dtype=jnp.floataa	)
         __lowerCamelCase   :	Optional[Any]								=		jnp.array(SCREAMING_SNAKE_CASE_			,     dtype=jnp.floataa	)

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







    @parameterized.expand(
        [
            # fmt: off
            [83, 4, [0.1_5_1_4, 0.0_8_0_7, 0.1_6_2_4, 0.1_0_1_6, -0.1_8_9_6, 0.0_2_6_3, 0.0_6_7_7, 0.2_3_1_0]],
            [17, 0.5_5, [0.1_1_6_4, -0.0_2_1_6, 0.0_1_7_0, 0.1_5_8_9, -0.3_1_2_0, 0.1_0_0_5, -0.0_5_8_1, -0.1_4_5_8]],
            [8, 0.8_9, [-0.1_7_5_8, -0.0_1_6_9, 0.1_0_0_4, -0.1_4_1_1, 0.1_3_1_2, 0.1_1_0_3, -0.1_9_9_6, 0.2_1_3_9]],
            [3, 10_00, [0.1_2_1_4, 0.0_3_5_2, -0.0_7_3_1, -0.1_5_6_2, -0.0_9_9_4, -0.0_9_0_6, -0.2_3_4_0, -0.0_5_3_9]],
            # fmt: on
        ]	)
    def       lowercase_				(							self			,     SCREAMING_SNAKE_CASE_			,     SCREAMING_SNAKE_CASE_			,     SCREAMING_SNAKE_CASE_	)				->     Optional[int]:
         __lowerCamelCase				,       __lowerCamelCase   :	List[str]								=		self.get_unet_model(model_id='stabilityai/stable-diffusion-2'			,     fpaa=SCREAMING_SNAKE_CASE_	)
         __lowerCamelCase   :	int								=		self.get_latents(SCREAMING_SNAKE_CASE_			,     shape=(4, 4, 96, 96)			,     fpaa=SCREAMING_SNAKE_CASE_	)
         __lowerCamelCase   :	Optional[int]								=		self.get_encoder_hidden_states(SCREAMING_SNAKE_CASE_			,     shape=(4, 77, 10_24)			,     fpaa=SCREAMING_SNAKE_CASE_	)

         __lowerCamelCase   :	List[str]								=		model.apply(
             {'params': params}			,     SCREAMING_SNAKE_CASE_			,     jnp.array(SCREAMING_SNAKE_CASE_			,     dtype=jnp.intaa	)			,     encoder_hidden_states=SCREAMING_SNAKE_CASE_			,     ).sample

         assert sample.shape == latents.shape

         __lowerCamelCase   :	Optional[int]								=		jnp.asarray(jax.device_get((sample[-1, -2:, -2:, :2].flatten())	)			,     dtype=jnp.floataa	)
         __lowerCamelCase   :	Tuple								=		jnp.array(SCREAMING_SNAKE_CASE_			,     dtype=jnp.floataa	)

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






