



import unittest

from transformers import is_torch_available
from transformers.testing_utils import require_torch


if is_torch_available():
							import torch

							from transformers.generation import DisjunctiveConstraint



@require_torch
class 	lowerCamelCase__      (							unittest.TestCase):




		'''simple docstring'''



		def 		lowerCAmelCase__				(self			)					->      Optional[int]:

						"""simple docstring"""






						lowerCAmelCase__					:    Optional[int]												=					[[1, 2, 4], [1, 2, 3, 4]]
						lowerCAmelCase__					:    Tuple												=					DisjunctiveConstraint(__lowerCamelCase			)
						self.assertTrue(isinstance(dc.token_ids		,__lowerCamelCase			)			)

						with self.assertRaises(__lowerCamelCase			):
										DisjunctiveConstraint(torch.LongTensor([[1, 2, 4], [1, 2, 3]]			)			)

						with self.assertRaises(__lowerCamelCase			):
										DisjunctiveConstraint([torch.LongTensor([1, 2, 4]			), torch.LongTensor([1, 2, 3, 4, 5]			)]			)



		def 		lowerCAmelCase__				(self			)					->      str:

						"""simple docstring"""






						lowerCAmelCase__					:    List[Any]												=					[[1, 2], [1, 2, 3, 4]]

						with self.assertRaises(__lowerCamelCase			):
										DisjunctiveConstraint(__lowerCamelCase			)  # fails here



		def 		lowerCAmelCase__				(self			)					->      Tuple:

						"""simple docstring"""






						lowerCAmelCase__					:    Any												=					[[1, 2, 3], [1, 2, 4]]

						lowerCAmelCase__					:    Dict												=					DisjunctiveConstraint(__lowerCamelCase			)

						lowerCAmelCase__						,		lowerCAmelCase__						,		lowerCAmelCase__					:    List[str]												=					dc.update(1			)
						lowerCAmelCase__					:    str												=					stepped is True and completed is False and reset is False
						self.assertTrue(__lowerCamelCase			)
						self.assertTrue(not dc.completed			)
						self.assertTrue(dc.current_seq == [1]			)

						lowerCAmelCase__						,		lowerCAmelCase__						,		lowerCAmelCase__					:    Any												=					dc.update(2			)
						lowerCAmelCase__					:    Union[str, Any]												=					stepped is True and completed is False and reset is False
						self.assertTrue(__lowerCamelCase			)
						self.assertTrue(not dc.completed			)
						self.assertTrue(dc.current_seq == [1, 2]			)

						lowerCAmelCase__						,		lowerCAmelCase__						,		lowerCAmelCase__					:    List[Any]												=					dc.update(3			)
						lowerCAmelCase__					:    List[str]												=					stepped is True and completed is True and reset is False
						self.assertTrue(__lowerCamelCase			)
						self.assertTrue(dc.completed			)  # Completed!
						self.assertTrue(dc.current_seq == [1, 2, 3]			)



		def 		lowerCAmelCase__				(self			)					->      List[str]:

						"""simple docstring"""






						lowerCAmelCase__					:    Optional[Any]												=					[[1, 2, 3], [1, 2, 4, 5], [1, 2, 5]]

						lowerCAmelCase__					:    Any												=					DisjunctiveConstraint(__lowerCamelCase			)

						lowerCAmelCase__						,		lowerCAmelCase__						,		lowerCAmelCase__					:    int												=					dc.update(1			)
						self.assertTrue(not dc.completed			)
						self.assertTrue(dc.current_seq == [1]			)

						lowerCAmelCase__						,		lowerCAmelCase__						,		lowerCAmelCase__					:    Union[str, Any]												=					dc.update(2			)
						self.assertTrue(not dc.completed			)
						self.assertTrue(dc.current_seq == [1, 2]			)

						lowerCAmelCase__						,		lowerCAmelCase__						,		lowerCAmelCase__					:    Tuple												=					dc.update(4			)
						self.assertTrue(not dc.completed			)
						self.assertTrue(dc.current_seq == [1, 2, 4]			)

						lowerCAmelCase__						,		lowerCAmelCase__						,		lowerCAmelCase__					:    Dict												=					dc.update(5			)
						self.assertTrue(dc.completed			)  # Completed!
						self.assertTrue(dc.current_seq == [1, 2, 4, 5]			)

						dc.reset()

						lowerCAmelCase__						,		lowerCAmelCase__						,		lowerCAmelCase__					:    Optional[Any]												=					dc.update(1			)
						self.assertTrue(not dc.completed			)
						self.assertTrue(dc.remaining() == 3			)
						self.assertTrue(dc.current_seq == [1]			)

						lowerCAmelCase__						,		lowerCAmelCase__						,		lowerCAmelCase__					:    Union[str, Any]												=					dc.update(2			)
						self.assertTrue(not dc.completed			)
						self.assertTrue(dc.remaining() == 2			)
						self.assertTrue(dc.current_seq == [1, 2]			)

						lowerCAmelCase__						,		lowerCAmelCase__						,		lowerCAmelCase__					:    Union[str, Any]												=					dc.update(5			)
						self.assertTrue(dc.completed			)  # Completed!
						self.assertTrue(dc.remaining() == 0			)
						self.assertTrue(dc.current_seq == [1, 2, 5]			)






import flax.linen as nn
import jax
import jax.numpy as jnp



class 	lowerCamelCase__      (							nn.Module):




		'''simple docstring'''



		snake_case_     =42
		snake_case_     =jnp.floataa



		def 		lowerCAmelCase__				(self			)					->      List[str]:

						"""simple docstring"""






						lowerCAmelCase__					:    Dict												=					nn.Conv(
						    self.out_channels		,kernel_size=(3, 3)		,strides=(1, 1)		,padding=((1, 1), (1, 1))		,dtype=self.dtype		,)



		def __call__(self		,__lowerCamelCase			)					->      Union[str, Any]:

						"""simple docstring"""






						lowerCAmelCase__						,		lowerCAmelCase__						,		lowerCAmelCase__						,		lowerCAmelCase__					:    Dict												=					hidden_states.shape
						lowerCAmelCase__					:    Dict												=					jax.image.resize(
						    __lowerCamelCase		,shape=(batch, height * 2, width * 2, channels)		,method='''nearest'''		,)
						lowerCAmelCase__					:    Dict												=					self.conv(__lowerCamelCase			)
						return hidden_states



class 	lowerCamelCase__      (							nn.Module):




		'''simple docstring'''



		snake_case_     =42
		snake_case_     =jnp.floataa



		def 		lowerCAmelCase__				(self			)					->      int:

						"""simple docstring"""






						lowerCAmelCase__					:    List[Any]												=					nn.Conv(
						    self.out_channels		,kernel_size=(3, 3)		,strides=(2, 2)		,padding=((1, 1), (1, 1))		,dtype=self.dtype		,)



		def __call__(self		,__lowerCamelCase			)					->      Any:

						"""simple docstring"""






						lowerCAmelCase__					:    Optional[int]												=					self.conv(__lowerCamelCase			)
						return hidden_states



class 	lowerCamelCase__      (							nn.Module):




		'''simple docstring'''



		snake_case_     =42
		snake_case_     =None
		snake_case_     =0.0
		snake_case_     =None
		snake_case_     =jnp.floataa



		def 		lowerCAmelCase__				(self			)					->      int:

						"""simple docstring"""






						lowerCAmelCase__					:    List[str]												=					self.in_channels if self.out_channels is None else self.out_channels

						lowerCAmelCase__					:    Union[str, Any]												=					nn.GroupNorm(num_groups=32		,epsilon=1e-5			)
						lowerCAmelCase__					:    Union[str, Any]												=					nn.Conv(
						    __lowerCamelCase		,kernel_size=(3, 3)		,strides=(1, 1)		,padding=((1, 1), (1, 1))		,dtype=self.dtype		,)

						lowerCAmelCase__					:    int												=					nn.Dense(__lowerCamelCase		,dtype=self.dtype			)

						lowerCAmelCase__					:    List[Any]												=					nn.GroupNorm(num_groups=32		,epsilon=1e-5			)
						lowerCAmelCase__					:    Union[str, Any]												=					nn.Dropout(self.dropout_prob			)
						lowerCAmelCase__					:    Optional[int]												=					nn.Conv(
						    __lowerCamelCase		,kernel_size=(3, 3)		,strides=(1, 1)		,padding=((1, 1), (1, 1))		,dtype=self.dtype		,)

						lowerCAmelCase__					:    str												=					self.in_channels != out_channels if self.use_nin_shortcut is None else self.use_nin_shortcut

						lowerCAmelCase__					:    Union[str, Any]												=					None
						if use_nin_shortcut:
										lowerCAmelCase__					:    Optional[Any]												=					nn.Conv(
										    __lowerCamelCase		,kernel_size=(1, 1)		,strides=(1, 1)		,padding='''VALID'''		,dtype=self.dtype		,)



		def __call__(self		,__lowerCamelCase		,__lowerCamelCase		,__lowerCamelCase=True			)					->      Union[str, Any]:

						"""simple docstring"""






						lowerCAmelCase__					:    Tuple												=					hidden_states
						lowerCAmelCase__					:    Dict												=					self.norma(__lowerCamelCase			)
						lowerCAmelCase__					:    Union[str, Any]												=					nn.swish(__lowerCamelCase			)
						lowerCAmelCase__					:    List[Any]												=					self.conva(__lowerCamelCase			)

						lowerCAmelCase__					:    List[Any]												=					self.time_emb_proj(nn.swish(__lowerCamelCase			)			)
						lowerCAmelCase__					:    Dict												=					jnp.expand_dims(jnp.expand_dims(__lowerCamelCase		,1			)		,1			)
						lowerCAmelCase__					:    Optional[int]												=					hidden_states + temb

						lowerCAmelCase__					:    Optional[int]												=					self.norma(__lowerCamelCase			)
						lowerCAmelCase__					:    Optional[int]												=					nn.swish(__lowerCamelCase			)
						lowerCAmelCase__					:    int												=					self.dropout(__lowerCamelCase		,__lowerCamelCase			)
						lowerCAmelCase__					:    List[Any]												=					self.conva(__lowerCamelCase			)

						if self.conv_shortcut is not None:
										lowerCAmelCase__					:    Optional[int]												=					self.conv_shortcut(__lowerCamelCase			)

						return hidden_states + residual

