





'''simple docstring'''




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







class   A_	(  nn.Module			):
 _lowerCamelCase  : int
 _lowerCamelCase  : jnp.dtype           =							jnp.floataa




 def      lowercase	(					self       :    int    ):
      _UpperCAmelCase     							=   nn.Conv(
          self.out_channels	,			kernel_size=(3, 3)	,			strides=(1, 1)	,			padding=((1, 1), (1, 1))	,			dtype=self.dtype	,			)




 def __call__(					self       :    Tuple	,			snake_case_       :    int    ):
      _UpperCAmelCase							,  _UpperCAmelCase							,  _UpperCAmelCase							,  _UpperCAmelCase     							=   hidden_states.shape
      _UpperCAmelCase     							=   jax.image.resize(
          snake_case_	,			shape=(batch, height * 2, width * 2, channels)	,			method="nearest"	,			)
      _UpperCAmelCase     							=   self.conv(snake_case_    )
      return hidden_states







class   A_	(  nn.Module			):
 _lowerCamelCase  : int
 _lowerCamelCase  : jnp.dtype           =							jnp.floataa




 def      lowercase	(					self       :    Dict    ):
      _UpperCAmelCase     							=   nn.Conv(
          self.out_channels	,			kernel_size=(3, 3)	,			strides=(2, 2)	,			padding=((1, 1), (1, 1))	,			dtype=self.dtype	,			)




 def __call__(					self       :    Tuple	,			snake_case_       :    Tuple    ):
      # 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(snake_case_    )
      return hidden_states







class   A_	(  nn.Module			):
 _lowerCamelCase  : int
 _lowerCamelCase  : int           =							None
 _lowerCamelCase  : float           =							0.0
 _lowerCamelCase  : bool           =							None
 _lowerCamelCase  : jnp.dtype           =							jnp.floataa




 def      lowercase	(					self       :    List[str]    ):
      _UpperCAmelCase     							=   self.in_channels if self.out_channels is None else self.out_channels

      _UpperCAmelCase     							=   nn.GroupNorm(num_groups=3_2	,			epsilon=1e-5    )
      _UpperCAmelCase     							=   nn.Conv(
          snake_case_	,			kernel_size=(3, 3)	,			strides=(1, 1)	,			padding=((1, 1), (1, 1))	,			dtype=self.dtype	,			)

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

      _UpperCAmelCase     							=   nn.GroupNorm(num_groups=3_2	,			epsilon=1e-5    )
      _UpperCAmelCase     							=   nn.Dropout(self.dropout_prob    )
      _UpperCAmelCase     							=   nn.Conv(
          snake_case_	,			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(
               snake_case_	,			kernel_size=(1, 1)	,			strides=(1, 1)	,			padding="VALID"	,			dtype=self.dtype	,			)




 def __call__(					self       :    Union[str, Any]	,			snake_case_       :    Optional[int]	,			snake_case_       :    Optional[Any]	,			snake_case_       :    List[Any]=True    ):
      _UpperCAmelCase     							=   hidden_states
      _UpperCAmelCase     							=   self.norma(snake_case_    )
      _UpperCAmelCase     							=   nn.swish(snake_case_    )
      _UpperCAmelCase     							=   self.conva(snake_case_    )

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

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

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

      return hidden_states + residual











'''simple docstring'''




def 	UpperCAmelCase_							(					__lowercase							:      int       )		->       int:


     '''simple docstring'''



     if not isinstance(__lowercase			,     __lowercase       ) or number < 0:
          raise ValueError("Input must be a non-negative integer"       )

     _UpperCAmelCase     							=   0
     while number:
          # This way we arrive at next set bit (next 1) instead of looping
          # through each bit and checking for 1s hence the
          # loop won't run 32 times it will only run the number of `1` times
          number &= number - 1
          count += 1
     return count


if __name__ == "__main__":
  import doctest

  doctest.testmod()




