


import os
import sys
import unittest


A  :     Optional[int]    					=    os.path.abspath(os.path.dirname(os.path.dirname(os.path.dirname(__file__))))
sys.path.append(os.path.join(git_repo_path, "utils"))

import check_dummies  # noqa: E402
from check_dummies import create_dummy_files, create_dummy_object, find_backend, read_init  # noqa: E402


# Align TRANSFORMERS_PATH in check_dummies with the current path
A  :     Optional[Any]    					=    os.path.join(git_repo_path, "src", "transformers")

A  :     str    					=    """
{0} = None
"""

A  :     int    					=    """
class {0}(metaclass=DummyObject):
    _backends = {1}

    def __init__(self, *args, **kwargs):
        requires_backends(self, {1})
"""


A  :     Any    					=    """
def {0}(*args, **kwargs):
    requires_backends({0}, {1})
"""
class 						_UpperCamelCase						(     unittest.TestCase	):
							'''simple docstring'''







							def     snake_case       (  self			):
										__lowerCAmelCase       					=    find_backend("    _import_structure[\"models.albert\"].append(\"AlbertTokenizerFast\")"			)
										self.assertIsNone(__a			)

										__lowerCAmelCase       					=    find_backend("    if not is_tokenizers_available():"			)
										self.assertEqual(__a					,     "tokenizers"			)

										__lowerCAmelCase       					=    find_backend("    if not is_tensorflow_text_available():"			)
										self.assertEqual(__a					,     "tensorflow_text"			)

										__lowerCAmelCase       					=    find_backend("    if not (is_sentencepiece_available() and is_tokenizers_available()):"			)
										self.assertEqual(__a					,     "sentencepiece_and_tokenizers"			)

										__lowerCAmelCase       					=    find_backend(
										    "    if not (is_sentencepiece_available() and is_tensorflow_text_available()):"			)
										self.assertEqual(__a					,     "sentencepiece_and_tensorflow_text"			)

										__lowerCAmelCase       					=    find_backend(
										    "    if not (is_sentencepiece_available() and is_tokenizers_available() and is_vision_available()):"			)
										self.assertEqual(__a					,     "sentencepiece_and_tokenizers_and_vision"			)

							def     snake_case       (  self			):
										__lowerCAmelCase       					=    read_init()
										# We don't assert on the exact list of keys to allow for smooth grow of backend-specific objects
										self.assertIn("torch"					,     __a			)
										self.assertIn("tensorflow_text"					,     __a			)
										self.assertIn("sentencepiece_and_tokenizers"					,     __a			)

										# Likewise, we can't assert on the exact content of a key
										self.assertIn("BertModel"					,     objects["torch"]			)
										self.assertIn("TFBertModel"					,     objects["tf"]			)
										self.assertIn("FlaxBertModel"					,     objects["flax"]			)
										self.assertIn("BertModel"					,     objects["torch"]			)
										self.assertIn("TFBertTokenizer"					,     objects["tensorflow_text"]			)
										self.assertIn("convert_slow_tokenizer"					,     objects["sentencepiece_and_tokenizers"]			)

							def     snake_case       (  self			):
										__lowerCAmelCase       					=    create_dummy_object("CONSTANT"					,     "\'torch\'"			)
										self.assertEqual(__a					,     "\nCONSTANT = None\n"			)

										__lowerCAmelCase       					=    create_dummy_object("function"					,     "\'torch\'"			)
										self.assertEqual(
										    __a					,     "\ndef function(*args, **kwargs):\n    requires_backends(function, \'torch\')\n"			)

										__lowerCAmelCase       					=    "\nclass FakeClass(metaclass=DummyObject):\n    _backends = \'torch\'\n\n    def __init__(self, *args, **kwargs):\n        requires_backends(self, \'torch\')\n"
										__lowerCAmelCase       					=    create_dummy_object("FakeClass"					,     "\'torch\'"			)
										self.assertEqual(__a					,     __a			)



							def     snake_case       (  self			):
										__lowerCAmelCase       					=    "# This file is autogenerated by the command `make fix-copies`, do not edit.\nfrom ..utils import DummyObject, requires_backends\n\n\nCONSTANT = None\n\n\ndef function(*args, **kwargs):\n    requires_backends(function, [\"torch\"])\n\n\nclass FakeClass(metaclass=DummyObject):\n    _backends = [\"torch\"]\n\n    def __init__(self, *args, **kwargs):\n        requires_backends(self, [\"torch\"])\n"
										__lowerCAmelCase       					=    create_dummy_files({"torch": ["CONSTANT", "function", "FakeClass"]}			)
										self.assertEqual(dummy_files["torch"]					,     __a			)








"""simple docstring"""


import flax.linen as nn
import jax
import jax.numpy as jnp
class 						_UpperCamelCase						(     nn.Module	):
							'''simple docstring'''







							__UpperCAmelCase		:					int
							__UpperCAmelCase		:					jnp.dtype          =jnp.floataa

							def     snake_case       (  self			):
										__lowerCAmelCase       					=    nn.Conv(
										    self.out_channels					,     kernel_size=(3, 3)					,     strides=(1, 1)					,     padding=((1, 1), (1, 1))					,     dtype=self.dtype					,     )



							def __call__(  self					,     __a			):
										__lowerCAmelCase					,			__lowerCAmelCase					,			__lowerCAmelCase					,			__lowerCAmelCase       					=    hidden_states.shape
										__lowerCAmelCase       					=    jax.image.resize(
										    __a					,     shape=(batch, height * 2, width * 2, channels)					,     method="nearest"					,     )
										__lowerCAmelCase       					=    self.conv(__a			)
										return hidden_states
class 						_UpperCamelCase						(     nn.Module	):
							'''simple docstring'''







							__UpperCAmelCase		:					int
							__UpperCAmelCase		:					jnp.dtype          =jnp.floataa

							def     snake_case       (  self			):
										__lowerCAmelCase       					=    nn.Conv(
										    self.out_channels					,     kernel_size=(3, 3)					,     strides=(2, 2)					,     padding=((1, 1), (1, 1))					,     dtype=self.dtype					,     )



							def __call__(  self					,     __a			):
										# pad = ((0, 0), (0, 1), (0, 1), (0, 0))  # pad height and width dim
										# hidden_states = jnp.pad(hidden_states, pad_width=pad)
										__lowerCAmelCase       					=    self.conv(__a			)
										return hidden_states
class 						_UpperCamelCase						(     nn.Module	):
							'''simple docstring'''







							__UpperCAmelCase		:					int
							__UpperCAmelCase		:					int          =None
							__UpperCAmelCase		:					float          =0.0
							__UpperCAmelCase		:					bool          =None
							__UpperCAmelCase		:					jnp.dtype          =jnp.floataa

							def     snake_case       (  self			):
										__lowerCAmelCase       					=    self.in_channels if self.out_channels is None else self.out_channels

										__lowerCAmelCase       					=    nn.GroupNorm(num_groups=32					,     epsilon=1e-5			)
										__lowerCAmelCase       					=    nn.Conv(
										    __a					,     kernel_size=(3, 3)					,     strides=(1, 1)					,     padding=((1, 1), (1, 1))					,     dtype=self.dtype					,     )

										__lowerCAmelCase       					=    nn.Dense(__a					,     dtype=self.dtype			)

										__lowerCAmelCase       					=    nn.GroupNorm(num_groups=32					,     epsilon=1e-5			)
										__lowerCAmelCase       					=    nn.Dropout(self.dropout_prob			)
										__lowerCAmelCase       					=    nn.Conv(
										    __a					,     kernel_size=(3, 3)					,     strides=(1, 1)					,     padding=((1, 1), (1, 1))					,     dtype=self.dtype					,     )

										__lowerCAmelCase       					=    self.in_channels != out_channels if self.use_nin_shortcut is None else self.use_nin_shortcut

										__lowerCAmelCase       					=    None
										if use_nin_shortcut:
													__lowerCAmelCase       					=    nn.Conv(
													    __a					,     kernel_size=(1, 1)					,     strides=(1, 1)					,     padding="VALID"					,     dtype=self.dtype					,     )



							def __call__(  self					,     __a					,     __a					,     __a=True			):
										__lowerCAmelCase       					=    hidden_states
										__lowerCAmelCase       					=    self.norma(__a			)
										__lowerCAmelCase       					=    nn.swish(__a			)
										__lowerCAmelCase       					=    self.conva(__a			)

										__lowerCAmelCase       					=    self.time_emb_proj(nn.swish(__a			)			)
										__lowerCAmelCase       					=    jnp.expand_dims(jnp.expand_dims(__a					,     1			)					,     1			)
										__lowerCAmelCase       					=    hidden_states + temb

										__lowerCAmelCase       					=    self.norma(__a			)
										__lowerCAmelCase       					=    nn.swish(__a			)
										__lowerCAmelCase       					=    self.dropout(__a					,     __a			)
										__lowerCAmelCase       					=    self.conva(__a			)

										if self.conv_shortcut is not None:
													__lowerCAmelCase       					=    self.conv_shortcut(__a			)

										return hidden_states + residual




