






import flax.linen as nn
import jax
import jax.numpy as jnp
class 		SCREAMING_SNAKE_CASE__	(  nn.Module     ):

  '''simple docstring'''







  _lowerCamelCase    =  42
  _lowerCamelCase    =  jnp.floataa


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


         """simple docstring"""
         _UpperCAmelCase							=    nn.Conv(
             self.out_channels       ,				kernel_size=(3, 3)       ,				strides=(1, 1)       ,				padding=((1, 1), (1, 1))       ,				dtype=self.dtype       ,				)





  def __call__(      self		:			int       ,				lowerCamelCase		:			List[Any]							)						->	List[str]:


         """simple docstring"""
         _UpperCAmelCase	,			_UpperCAmelCase	,			_UpperCAmelCase	,			_UpperCAmelCase							=    hidden_states.shape
         _UpperCAmelCase							=    jax.image.resize(
             lowerCamelCase       ,				shape=(batch, height * 2, width * 2, channels)       ,				method="""nearest"""       ,				)
         _UpperCAmelCase							=    self.conv(lowerCamelCase							)
         return hidden_states



class 		SCREAMING_SNAKE_CASE__	(  nn.Module     ):

  '''simple docstring'''







  _lowerCamelCase    =  42
  _lowerCamelCase    =  jnp.floataa


  def  lowerCamelCase       (      self		:			Any							)						->	int:


         """simple docstring"""
         _UpperCAmelCase							=    nn.Conv(
             self.out_channels       ,				kernel_size=(3, 3)       ,				strides=(2, 2)       ,				padding=((1, 1), (1, 1))       ,				dtype=self.dtype       ,				)





  def __call__(      self		:			Any       ,				lowerCamelCase		:			List[Any]							)						->	Optional[Any]:


         """simple docstring"""
         # pad = ((0, 0), (0, 1), (0, 1), (0, 0))  # pad height and width dim
         # hidden_states = jnp.pad(hidden_states, pad_width=pad)
         _UpperCAmelCase							=    self.conv(lowerCamelCase							)
         return hidden_states



class 		SCREAMING_SNAKE_CASE__	(  nn.Module     ):

  '''simple docstring'''







  _lowerCamelCase    =  42
  _lowerCamelCase    =  None
  _lowerCamelCase    =  0.0
  _lowerCamelCase    =  None
  _lowerCamelCase    =  jnp.floataa


  def  lowerCamelCase       (      self		:			Any							)						->	List[str]:


         """simple docstring"""
         _UpperCAmelCase							=    self.in_channels if self.out_channels is None else self.out_channels

         _UpperCAmelCase							=    nn.GroupNorm(num_groups=32       ,				epsilon=1E-5							)
         _UpperCAmelCase							=    nn.Conv(
             lowerCamelCase       ,				kernel_size=(3, 3)       ,				strides=(1, 1)       ,				padding=((1, 1), (1, 1))       ,				dtype=self.dtype       ,				)

         _UpperCAmelCase							=    nn.Dense(lowerCamelCase       ,				dtype=self.dtype							)

         _UpperCAmelCase							=    nn.GroupNorm(num_groups=32       ,				epsilon=1E-5							)
         _UpperCAmelCase							=    nn.Dropout(self.dropout_prob							)
         _UpperCAmelCase							=    nn.Conv(
             lowerCamelCase       ,				kernel_size=(3, 3)       ,				strides=(1, 1)       ,				padding=((1, 1), (1, 1))       ,				dtype=self.dtype       ,				)

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

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





  def __call__(      self		:			Dict       ,				lowerCamelCase		:			int       ,				lowerCamelCase		:			int       ,				lowerCamelCase		:			Tuple=True							)						->	Any:


         """simple docstring"""
         _UpperCAmelCase							=    hidden_states
         _UpperCAmelCase							=    self.norma(lowerCamelCase							)
         _UpperCAmelCase							=    nn.swish(lowerCamelCase							)
         _UpperCAmelCase							=    self.conva(lowerCamelCase							)

         _UpperCAmelCase							=    self.time_emb_proj(nn.swish(lowerCamelCase							)							)
         _UpperCAmelCase							=    jnp.expand_dims(jnp.expand_dims(lowerCamelCase       ,				1							)       ,				1							)
         _UpperCAmelCase							=    hidden_states + temb

         _UpperCAmelCase							=    self.norma(lowerCamelCase							)
         _UpperCAmelCase							=    nn.swish(lowerCamelCase							)
         _UpperCAmelCase							=    self.dropout(lowerCamelCase       ,				lowerCamelCase							)
         _UpperCAmelCase							=    self.conva(lowerCamelCase							)

         if self.conv_shortcut is not None:
                _UpperCAmelCase							=    self.conv_shortcut(lowerCamelCase							)

         return hidden_states + residual







# Copyright 2023 The HuggingFace Team. All rights reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

from typing import TYPE_CHECKING

from ...utils import OptionalDependencyNotAvailable, _LazyModule, is_torch_available


_SCREAMING_SNAKE_CASE   =  {
    "configuration_xmod": [
        "XMOD_PRETRAINED_CONFIG_ARCHIVE_MAP",
        "XmodConfig",
        "XmodOnnxConfig",
    ],
}

try:
    if not is_torch_available():
        raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
    pass
else:
    _SCREAMING_SNAKE_CASE   =  [
        "XMOD_PRETRAINED_MODEL_ARCHIVE_LIST",
        "XmodForCausalLM",
        "XmodForMaskedLM",
        "XmodForMultipleChoice",
        "XmodForQuestionAnswering",
        "XmodForSequenceClassification",
        "XmodForTokenClassification",
        "XmodModel",
        "XmodPreTrainedModel",
    ]

if TYPE_CHECKING:
    from .configuration_xmod import XMOD_PRETRAINED_CONFIG_ARCHIVE_MAP, XmodConfig, XmodOnnxConfig

    try:
        if not is_torch_available():
            raise OptionalDependencyNotAvailable()
    except OptionalDependencyNotAvailable:
        pass
    else:
        from .modeling_xmod import (
            XMOD_PRETRAINED_MODEL_ARCHIVE_LIST,
            XmodForCausalLM,
            XmodForMaskedLM,
            XmodForMultipleChoice,
            XmodForQuestionAnswering,
            XmodForSequenceClassification,
            XmodForTokenClassification,
            XmodModel,
            XmodPreTrainedModel,
        )

else:
    import sys

    _SCREAMING_SNAKE_CASE   =  _LazyModule(__name__, globals()["__file__"], _import_structure, module_spec=__spec__)


