


'''simple docstring'''






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

from .attention_flax import FlaxTransformeraDModel
from .resnet_flax import FlaxDownsampleaD, FlaxResnetBlockaD, FlaxUpsampleaD






class  __magic_name__							(		nn.Module    ):
       UpperCAmelCase										=4_2
       UpperCAmelCase										=4_2
       UpperCAmelCase										=0.0
       UpperCAmelCase										=1
       UpperCAmelCase										=1
       UpperCAmelCase										=True
       UpperCAmelCase										=False
       UpperCAmelCase										=False
       UpperCAmelCase										=False
       UpperCAmelCase										=jnp.floataa






       def       lowerCAmelCase		(      self)						->  Tuple:




          '''simple docstring'''
          _UpperCAmelCase       :       Tuple							   =[]
          _UpperCAmelCase       :       int							   =[]

          for i in range(self.num_layers):
             _UpperCAmelCase       :       Optional[Any]							   =self.in_channels if i == 0 else self.out_channels

             _UpperCAmelCase       :       Optional[int]							   =FlaxResnetBlockaD(
                 in_channels=snake_case  ,							out_channels=self.out_channels  ,							dropout_prob=self.dropout  ,							dtype=self.dtype  ,							)
             resnets.append(snake_case)

             _UpperCAmelCase       :       Union[str, Any]							   =FlaxTransformeraDModel(
                 in_channels=self.out_channels  ,							n_heads=self.num_attention_heads  ,							d_head=self.out_channels // self.num_attention_heads  ,							depth=1  ,							use_linear_projection=self.use_linear_projection  ,							only_cross_attention=self.only_cross_attention  ,							use_memory_efficient_attention=self.use_memory_efficient_attention  ,							dtype=self.dtype  ,							)
             attentions.append(snake_case)

          _UpperCAmelCase       :       str							   =resnets
          _UpperCAmelCase       :       Optional[int]							   =attentions

          if self.add_downsample:
             _UpperCAmelCase       :       Any							   =FlaxDownsampleaD(self.out_channels  ,							dtype=self.dtype)






       def __call__(      self  ,							snake_case  ,							snake_case  ,							snake_case  ,							snake_case=True)						->  Optional[int]:




          '''simple docstring'''
          _UpperCAmelCase       :       int							   =()

          for resnet, attn in zip(self.resnets  ,							self.attentions):
             _UpperCAmelCase       :       int							   =resnet(snake_case  ,							snake_case  ,							deterministic=snake_case)
             _UpperCAmelCase       :       Tuple							   =attn(snake_case  ,							snake_case  ,							deterministic=snake_case)
             output_states += (hidden_states,)

          if self.add_downsample:
             _UpperCAmelCase       :       Optional[Any]							   =self.downsamplers_a(snake_case)
             output_states += (hidden_states,)

          return hidden_states, output_states






class  __magic_name__							(		nn.Module    ):
       UpperCAmelCase										=4_2
       UpperCAmelCase										=4_2
       UpperCAmelCase										=0.0
       UpperCAmelCase										=1
       UpperCAmelCase										=True
       UpperCAmelCase										=jnp.floataa






       def       lowerCAmelCase		(      self)						->  str:




          '''simple docstring'''
          _UpperCAmelCase       :       str							   =[]

          for i in range(self.num_layers):
             _UpperCAmelCase       :       Any							   =self.in_channels if i == 0 else self.out_channels

             _UpperCAmelCase       :       Tuple							   =FlaxResnetBlockaD(
                 in_channels=snake_case  ,							out_channels=self.out_channels  ,							dropout_prob=self.dropout  ,							dtype=self.dtype  ,							)
             resnets.append(snake_case)
          _UpperCAmelCase       :       Union[str, Any]							   =resnets

          if self.add_downsample:
             _UpperCAmelCase       :       Any							   =FlaxDownsampleaD(self.out_channels  ,							dtype=self.dtype)






       def __call__(      self  ,							snake_case  ,							snake_case  ,							snake_case=True)						->  Tuple:




          '''simple docstring'''
          _UpperCAmelCase       :       Tuple							   =()

          for resnet in self.resnets:
             _UpperCAmelCase       :       str							   =resnet(snake_case  ,							snake_case  ,							deterministic=snake_case)
             output_states += (hidden_states,)

          if self.add_downsample:
             _UpperCAmelCase       :       List[Any]							   =self.downsamplers_a(snake_case)
             output_states += (hidden_states,)

          return hidden_states, output_states






class  __magic_name__							(		nn.Module    ):
       UpperCAmelCase										=4_2
       UpperCAmelCase										=4_2
       UpperCAmelCase										=4_2
       UpperCAmelCase										=0.0
       UpperCAmelCase										=1
       UpperCAmelCase										=1
       UpperCAmelCase										=True
       UpperCAmelCase										=False
       UpperCAmelCase										=False
       UpperCAmelCase										=False
       UpperCAmelCase										=jnp.floataa






       def       lowerCAmelCase		(      self)						->  str:




          '''simple docstring'''
          _UpperCAmelCase       :       List[str]							   =[]
          _UpperCAmelCase       :       Optional[Any]							   =[]

          for i in range(self.num_layers):
             _UpperCAmelCase       :       int							   =self.in_channels if (i == self.num_layers - 1) else self.out_channels
             _UpperCAmelCase       :       Dict							   =self.prev_output_channel if i == 0 else self.out_channels

             _UpperCAmelCase       :       List[Any]							   =FlaxResnetBlockaD(
                 in_channels=resnet_in_channels + res_skip_channels  ,							out_channels=self.out_channels  ,							dropout_prob=self.dropout  ,							dtype=self.dtype  ,							)
             resnets.append(snake_case)

             _UpperCAmelCase       :       Optional[int]							   =FlaxTransformeraDModel(
                 in_channels=self.out_channels  ,							n_heads=self.num_attention_heads  ,							d_head=self.out_channels // self.num_attention_heads  ,							depth=1  ,							use_linear_projection=self.use_linear_projection  ,							only_cross_attention=self.only_cross_attention  ,							use_memory_efficient_attention=self.use_memory_efficient_attention  ,							dtype=self.dtype  ,							)
             attentions.append(snake_case)

          _UpperCAmelCase       :       Optional[Any]							   =resnets
          _UpperCAmelCase       :       Any							   =attentions

          if self.add_upsample:
             _UpperCAmelCase       :       Any							   =FlaxUpsampleaD(self.out_channels  ,							dtype=self.dtype)






       def __call__(      self  ,							snake_case  ,							snake_case  ,							snake_case  ,							snake_case  ,							snake_case=True)						->  Optional[int]:




          '''simple docstring'''
          for resnet, attn in zip(self.resnets  ,							self.attentions):
             # pop res hidden states
             _UpperCAmelCase       :       Union[str, Any]							   =res_hidden_states_tuple[-1]
             _UpperCAmelCase       :       Any							   =res_hidden_states_tuple[:-1]
             _UpperCAmelCase       :       str							   =jnp.concatenate((hidden_states, res_hidden_states)  ,							axis=-1)

             _UpperCAmelCase       :       Union[str, Any]							   =resnet(snake_case  ,							snake_case  ,							deterministic=snake_case)
             _UpperCAmelCase       :       Dict							   =attn(snake_case  ,							snake_case  ,							deterministic=snake_case)

          if self.add_upsample:
             _UpperCAmelCase       :       List[Any]							   =self.upsamplers_a(snake_case)

          return hidden_states






class  __magic_name__							(		nn.Module    ):
       UpperCAmelCase										=4_2
       UpperCAmelCase										=4_2
       UpperCAmelCase										=4_2
       UpperCAmelCase										=0.0
       UpperCAmelCase										=1
       UpperCAmelCase										=True
       UpperCAmelCase										=jnp.floataa






       def       lowerCAmelCase		(      self)						->  Any:




          '''simple docstring'''
          _UpperCAmelCase       :       Optional[Any]							   =[]

          for i in range(self.num_layers):
             _UpperCAmelCase       :       List[Any]							   =self.in_channels if (i == self.num_layers - 1) else self.out_channels
             _UpperCAmelCase       :       Tuple							   =self.prev_output_channel if i == 0 else self.out_channels

             _UpperCAmelCase       :       Any							   =FlaxResnetBlockaD(
                 in_channels=resnet_in_channels + res_skip_channels  ,							out_channels=self.out_channels  ,							dropout_prob=self.dropout  ,							dtype=self.dtype  ,							)
             resnets.append(snake_case)

          _UpperCAmelCase       :       Optional[Any]							   =resnets

          if self.add_upsample:
             _UpperCAmelCase       :       int							   =FlaxUpsampleaD(self.out_channels  ,							dtype=self.dtype)






       def __call__(      self  ,							snake_case  ,							snake_case  ,							snake_case  ,							snake_case=True)						->  Tuple:




          '''simple docstring'''
          for resnet in self.resnets:
             # pop res hidden states
             _UpperCAmelCase       :       Dict							   =res_hidden_states_tuple[-1]
             _UpperCAmelCase       :       Any							   =res_hidden_states_tuple[:-1]
             _UpperCAmelCase       :       str							   =jnp.concatenate((hidden_states, res_hidden_states)  ,							axis=-1)

             _UpperCAmelCase       :       Optional[int]							   =resnet(snake_case  ,							snake_case  ,							deterministic=snake_case)

          if self.add_upsample:
             _UpperCAmelCase       :       Dict							   =self.upsamplers_a(snake_case)

          return hidden_states






class  __magic_name__							(		nn.Module    ):
       UpperCAmelCase										=4_2
       UpperCAmelCase										=0.0
       UpperCAmelCase										=1
       UpperCAmelCase										=1
       UpperCAmelCase										=False
       UpperCAmelCase										=False
       UpperCAmelCase										=jnp.floataa






       def       lowerCAmelCase		(      self)						->  List[Any]:




          '''simple docstring'''
          _UpperCAmelCase       :       str							   =[
              FlaxResnetBlockaD(
                  in_channels=self.in_channels  ,							out_channels=self.in_channels  ,							dropout_prob=self.dropout  ,							dtype=self.dtype  ,							)
          ]

          _UpperCAmelCase       :       Tuple							   =[]

          for _ in range(self.num_layers):
             _UpperCAmelCase       :       Any							   =FlaxTransformeraDModel(
                 in_channels=self.in_channels  ,							n_heads=self.num_attention_heads  ,							d_head=self.in_channels // self.num_attention_heads  ,							depth=1  ,							use_linear_projection=self.use_linear_projection  ,							use_memory_efficient_attention=self.use_memory_efficient_attention  ,							dtype=self.dtype  ,							)
             attentions.append(snake_case)

             _UpperCAmelCase       :       Optional[int]							   =FlaxResnetBlockaD(
                 in_channels=self.in_channels  ,							out_channels=self.in_channels  ,							dropout_prob=self.dropout  ,							dtype=self.dtype  ,							)
             resnets.append(snake_case)

          _UpperCAmelCase       :       Optional[Any]							   =resnets
          _UpperCAmelCase       :       int							   =attentions






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




          '''simple docstring'''
          _UpperCAmelCase       :       List[Any]							   =self.resnets[0](snake_case  ,							snake_case)
          for attn, resnet in zip(self.attentions  ,							self.resnets[1:]):
             _UpperCAmelCase       :       Tuple							   =attn(snake_case  ,							snake_case  ,							deterministic=snake_case)
             _UpperCAmelCase       :       List[str]							   =resnet(snake_case  ,							snake_case  ,							deterministic=snake_case)

          return hidden_states





'''simple docstring'''






from string import ascii_uppercase

lowercase							={str(ord(c) - 55): c for c in ascii_uppercase}





def lowerCamelCase__      (						__lowerCamelCase	:							int	,   __lowerCamelCase	:							int							):







   '''simple docstring'''




   if isinstance(__lowerCamelCase	,   __lowerCamelCase							):
      raise TypeError('int() can\'t convert non-string with explicit base'							)
   if num < 0:
      raise ValueError('parameter must be positive int'							)
   if isinstance(__lowerCamelCase	,   __lowerCamelCase							):
      raise TypeError('\'str\' object cannot be interpreted as an integer'							)
   if isinstance(__lowerCamelCase	,   __lowerCamelCase							):
      raise TypeError('\'float\' object cannot be interpreted as an integer'							)
   if base in (0, 1):
      raise ValueError('base must be >= 2'							)
   if base > 3_6:
      raise ValueError('base must be <= 36'							)
   _UpperCAmelCase       :       Union[str, Any]							   =''
   _UpperCAmelCase       :       Optional[int]							   =0
   _UpperCAmelCase       :       str							   =0
   while div != 1:
      _UpperCAmelCase   ,	_UpperCAmelCase       :       int							   =divmod(__lowerCamelCase	,   __lowerCamelCase							)
      if base >= 1_1 and 9 < mod < 3_6:
         _UpperCAmelCase       :       str							   =ALPHABET_VALUES[str(__lowerCamelCase							)]
      else:
         _UpperCAmelCase       :       Any							   =str(__lowerCamelCase							)
      new_value += actual_value
      _UpperCAmelCase       :       Union[str, Any]							   =num // base
      _UpperCAmelCase       :       Dict							   =div
      if div == 0:
         return str(new_value[::-1]							)
      elif div == 1:
         new_value += str(__lowerCamelCase							)
         return str(new_value[::-1]							)

   return new_value[::-1]


if __name__ == "__main__":
 import doctest

 doctest.testmod()

 for base in range(2, 37):
  for num in range(1000):
   assert int(decimal_to_any(num, base), base) == num, (
       num,
       base,
       decimal_to_any(num, base),
       int(decimal_to_any(num, base), base),
   )

