




"""simple docstring"""







from dataclasses import dataclass
from typing import Optional

import torch
from torch import nn

from ..configuration_utils import ConfigMixin, register_to_config
from ..utils import BaseOutput
from .attention import BasicTransformerBlock
from .modeling_utils import ModelMixin




@dataclass
class       lowerCAmelCase__       (			lowerCamelCase__    ):


				"""simple docstring"""





				__UpperCAmelCase :     Optional[Any]       =      42





class       lowerCAmelCase__       (			lowerCamelCase__  ,			lowerCamelCase__    ):


				"""simple docstring"""





				@register_to_config
				def __init__(     self					:	List[str]    ,   lowercase__					:	List[str] = 1_6    ,   lowercase__					:	Any = 8_8    ,   lowercase__					:	Tuple = None    ,   lowercase__					:	Optional[int] = None    ,   lowercase__					:	List[str] = 1    ,   lowercase__					:	Optional[Any] = 0.0    ,   lowercase__					:	List[Any] = 3_2    ,   lowercase__					:	Union[str, Any] = None    ,   lowercase__					:	Any = False    ,   lowercase__					:	str = None    ,   lowercase__					:	Optional[Any] = "geglu"    ,   lowercase__					:	List[Any] = True    ,   lowercase__					:	Dict = True    ,   ):
						super().__init__()
						__lowercase				:		List[str]												=				num_attention_heads
						__lowercase				:		int												=				attention_head_dim
						__lowercase				:		Any												=				num_attention_heads * attention_head_dim

						__lowercase				:		Optional[Any]												=				in_channels

						__lowercase				:		Optional[Any]												=				torch.nn.GroupNorm(num_groups=__lowerCamelCase    ,   num_channels=__lowerCamelCase    ,   eps=1e-6    ,   affine=__lowerCamelCase						)
						__lowercase				:		Optional[Any]												=				nn.Linear(__lowerCamelCase    ,   __lowerCamelCase						)

						# 3. Define transformers blocks
						__lowercase				:		List[str]												=				nn.ModuleList(
						    [
						        BasicTransformerBlock(
						            __lowerCamelCase    ,   __lowerCamelCase    ,   __lowerCamelCase    ,   dropout=__lowerCamelCase    ,   cross_attention_dim=__lowerCamelCase    ,   activation_fn=__lowerCamelCase    ,   attention_bias=__lowerCamelCase    ,   double_self_attention=__lowerCamelCase    ,   norm_elementwise_affine=__lowerCamelCase    ,   )
						        for d in range(__lowerCamelCase						)
						    ]						)

						__lowercase				:		List[Any]												=				nn.Linear(__lowerCamelCase    ,   __lowerCamelCase						)



				def       snake_case    (     self					:	List[Any]    ,   lowercase__					:	Tuple    ,   lowercase__					:	List[str]=None    ,   lowercase__					:	Dict=None    ,   lowercase__					:	Dict=None    ,   lowercase__					:	Union[str, Any]=1    ,   lowercase__					:	List[str]=None    ,   lowercase__					:	List[str] = True    ,   ):
						__lowercase				:		str												=				hidden_states.shape
						__lowercase				:		Optional[int]												=				batch_frames // num_frames

						__lowercase				:		Optional[int]												=				hidden_states

						__lowercase				:		str												=				hidden_states[None, :].reshape(__lowerCamelCase    ,   __lowerCamelCase    ,   __lowerCamelCase    ,   __lowerCamelCase    ,   __lowerCamelCase						)
						__lowercase				:		List[Any]												=				hidden_states.permute(0    ,   2    ,   1    ,   3    ,   4						)

						__lowercase				:		str												=				self.norm(__lowerCamelCase						)
						__lowercase				:		Any												=				hidden_states.permute(0    ,   3    ,   4    ,   2    ,   1						).reshape(batch_size * height * width    ,   __lowerCamelCase    ,   __lowerCamelCase						)

						__lowercase				:		Optional[int]												=				self.proj_in(__lowerCamelCase						)

						# 2. Blocks
						for block in self.transformer_blocks:
								__lowercase				:		Union[str, Any]												=				block(
								    __lowerCamelCase    ,   encoder_hidden_states=__lowerCamelCase    ,   timestep=__lowerCamelCase    ,   cross_attention_kwargs=__lowerCamelCase    ,   class_labels=__lowerCamelCase    ,   )

						# 3. Output
						__lowercase				:		str												=				self.proj_out(__lowerCamelCase						)
						__lowercase				:		List[Any]												=				(
						    hidden_states[None, None, :]
						    .reshape(__lowerCamelCase    ,   __lowerCamelCase    ,   __lowerCamelCase    ,   __lowerCamelCase    ,   __lowerCamelCase						)
						    .permute(0    ,   3    ,   4    ,   1    ,   2						)
						    .contiguous()
						)
						__lowercase				:		Optional[int]												=				hidden_states.reshape(__lowerCamelCase    ,   __lowerCamelCase    ,   __lowerCamelCase    ,   __lowerCamelCase						)

						__lowercase				:		int												=				hidden_states + residual

						if not return_dict:
								return (output,)

						return TransformerTemporalModelOutput(sample=__lowerCamelCase						)






"""simple docstring"""







import unittest

from transformers.testing_utils import CaptureStdout
from transformers.tools.python_interpreter import evaluate

def 						snake_case__       (					_lowerCamelCase				)			->Dict:






		"""simple docstring"""


		return x + 2





class       lowerCAmelCase__       (			unittest.TestCase    ):


				"""simple docstring"""





				def       snake_case    (     self					:	List[str]						):
						__lowercase				:		List[Any]												=				"x = 3"
						__lowercase				:		Optional[int]												=				{}
						__lowercase				:		List[Any]												=				evaluate(lowercase__    ,   {}    ,   state=lowercase__						)
						assert result == 3
						self.assertDictEqual(lowercase__    ,   {"x": 3}						)

						__lowercase				:		Optional[int]												=				"x = y"
						__lowercase				:		Union[str, Any]												=				{"y": 5}
						__lowercase				:		str												=				evaluate(lowercase__    ,   {}    ,   state=lowercase__						)
						# evaluate returns the value of the last assignment.
						assert result == 5
						self.assertDictEqual(lowercase__    ,   {"x": 5, "y": 5}						)



				def       snake_case    (     self					:	Dict						):
						__lowercase				:		Dict												=				"y = add_two(x)"
						__lowercase				:		List[Any]												=				{"x": 3}
						__lowercase				:		Optional[int]												=				evaluate(lowercase__    ,   {"add_two": add_two}    ,   state=lowercase__						)
						assert result == 5
						self.assertDictEqual(lowercase__    ,   {"x": 3, "y": 5}						)

						# Won't work without the tool
						with CaptureStdout() as out:
								__lowercase				:		Any												=				evaluate(lowercase__    ,   {}    ,   state=lowercase__						)
						assert result is None
						assert "tried to execute add_two" in out.out



				def       snake_case    (     self					:	Any						):
						__lowercase				:		Optional[Any]												=				"x = 3"
						__lowercase				:		List[Any]												=				{}
						__lowercase				:		Dict												=				evaluate(lowercase__    ,   {}    ,   state=lowercase__						)
						assert result == 3
						self.assertDictEqual(lowercase__    ,   {"x": 3}						)



				def       snake_case    (     self					:	str						):
						__lowercase				:		Optional[int]												=				"test_dict = {'x': x, 'y': add_two(x)}"
						__lowercase				:		List[str]												=				{"x": 3}
						__lowercase				:		Tuple												=				evaluate(lowercase__    ,   {"add_two": add_two}    ,   state=lowercase__						)
						self.assertDictEqual(lowercase__    ,   {"x": 3, "y": 5}						)
						self.assertDictEqual(lowercase__    ,   {"x": 3, "test_dict": {"x": 3, "y": 5}}						)



				def       snake_case    (     self					:	int						):
						__lowercase				:		Tuple												=				"x = 3\ny = 5"
						__lowercase				:		Optional[Any]												=				{}
						__lowercase				:		List[Any]												=				evaluate(lowercase__    ,   {}    ,   state=lowercase__						)
						# evaluate returns the value of the last assignment.
						assert result == 5
						self.assertDictEqual(lowercase__    ,   {"x": 3, "y": 5}						)



				def       snake_case    (     self					:	Any						):
						__lowercase				:		str												=				"text = f'This is x: {x}.'"
						__lowercase				:		Any												=				{"x": 3}
						__lowercase				:		List[Any]												=				evaluate(lowercase__    ,   {}    ,   state=lowercase__						)
						# evaluate returns the value of the last assignment.
						assert result == "This is x: 3."
						self.assertDictEqual(lowercase__    ,   {"x": 3, "text": "This is x: 3."}						)



				def       snake_case    (     self					:	str						):
						__lowercase				:		Optional[Any]												=				"if x <= 3:\n    y = 2\nelse:\n    y = 5"
						__lowercase				:		int												=				{"x": 3}
						__lowercase				:		Optional[int]												=				evaluate(lowercase__    ,   {}    ,   state=lowercase__						)
						# evaluate returns the value of the last assignment.
						assert result == 2
						self.assertDictEqual(lowercase__    ,   {"x": 3, "y": 2}						)

						__lowercase				:		Optional[int]												=				{"x": 8}
						__lowercase				:		Union[str, Any]												=				evaluate(lowercase__    ,   {}    ,   state=lowercase__						)
						# evaluate returns the value of the last assignment.
						assert result == 5
						self.assertDictEqual(lowercase__    ,   {"x": 8, "y": 5}						)



				def       snake_case    (     self					:	Dict						):
						__lowercase				:		Optional[Any]												=				"test_list = [x, add_two(x)]"
						__lowercase				:		str												=				{"x": 3}
						__lowercase				:		Any												=				evaluate(lowercase__    ,   {"add_two": add_two}    ,   state=lowercase__						)
						self.assertListEqual(lowercase__    ,   [3, 5]						)
						self.assertDictEqual(lowercase__    ,   {"x": 3, "test_list": [3, 5]}						)



				def       snake_case    (     self					:	Optional[int]						):
						__lowercase				:		str												=				"y = x"
						__lowercase				:		str												=				{"x": 3}
						__lowercase				:		int												=				evaluate(lowercase__    ,   {}    ,   state=lowercase__						)
						assert result == 3
						self.assertDictEqual(lowercase__    ,   {"x": 3, "y": 3}						)



				def       snake_case    (     self					:	Optional[int]						):
						__lowercase				:		str												=				"test_list = [x, add_two(x)]\ntest_list[1]"
						__lowercase				:		Optional[Any]												=				{"x": 3}
						__lowercase				:		Any												=				evaluate(lowercase__    ,   {"add_two": add_two}    ,   state=lowercase__						)
						assert result == 5
						self.assertDictEqual(lowercase__    ,   {"x": 3, "test_list": [3, 5]}						)

						__lowercase				:		Tuple												=				"test_dict = {'x': x, 'y': add_two(x)}\ntest_dict['y']"
						__lowercase				:		Union[str, Any]												=				{"x": 3}
						__lowercase				:		Tuple												=				evaluate(lowercase__    ,   {"add_two": add_two}    ,   state=lowercase__						)
						assert result == 5
						self.assertDictEqual(lowercase__    ,   {"x": 3, "test_dict": {"x": 3, "y": 5}}						)



				def       snake_case    (     self					:	Dict						):
						__lowercase				:		Dict												=				"x = 0\nfor i in range(3):\n    x = i"
						__lowercase				:		str												=				{}
						__lowercase				:		int												=				evaluate(lowercase__    ,   {"range": range}    ,   state=lowercase__						)
						assert result == 2
						self.assertDictEqual(lowercase__    ,   {"x": 2, "i": 2}						)
