

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



class       __snake_case  (    nn.Module       ):
  __lowerCAmelCase     :	int
  __lowerCAmelCase     :	jnp.dtype						 =  jnp.floataa

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


  def __call__(					self						,		_A):
      SCREAMING_SNAKE_CASE_							,			SCREAMING_SNAKE_CASE_							,			SCREAMING_SNAKE_CASE_							,			SCREAMING_SNAKE_CASE_									=						hidden_states.shape
      SCREAMING_SNAKE_CASE_									=						jax.image.resize(
          _A						,		shape=(batch, height * 2, width * 2, channels)						,		method='nearest'						,		)
      SCREAMING_SNAKE_CASE_									=						self.conv(_A)
      return hidden_states






class       __snake_case  (    nn.Module       ):
  __lowerCAmelCase     :	int
  __lowerCAmelCase     :	jnp.dtype						 =  jnp.floataa

  def 						lowerCAmelCase__     (					self):
      SCREAMING_SNAKE_CASE_									=						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)
      SCREAMING_SNAKE_CASE_									=						self.conv(_A)
      return hidden_states






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

  def 						lowerCAmelCase__     (					self):
      SCREAMING_SNAKE_CASE_									=						self.in_channels if self.out_channels is None else self.out_channels

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

      SCREAMING_SNAKE_CASE_									=						nn.Dense(_A						,		dtype=self.dtype)

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

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

      SCREAMING_SNAKE_CASE_									=						None
      if use_nin_shortcut:
          SCREAMING_SNAKE_CASE_									=						nn.Conv(
              _A						,		kernel_size=(1, 1)						,		strides=(1, 1)						,		padding='VALID'						,		dtype=self.dtype						,		)


  def __call__(					self						,		_A						,		_A						,		_A=True):
      SCREAMING_SNAKE_CASE_									=						hidden_states
      SCREAMING_SNAKE_CASE_									=						self.norma(_A)
      SCREAMING_SNAKE_CASE_									=						nn.swish(_A)
      SCREAMING_SNAKE_CASE_									=						self.conva(_A)

      SCREAMING_SNAKE_CASE_									=						self.time_emb_proj(nn.swish(_A))
      SCREAMING_SNAKE_CASE_									=						jnp.expand_dims(jnp.expand_dims(_A						,		1)						,		1)
      SCREAMING_SNAKE_CASE_									=						hidden_states + temb

      SCREAMING_SNAKE_CASE_									=						self.norma(_A)
      SCREAMING_SNAKE_CASE_									=						nn.swish(_A)
      SCREAMING_SNAKE_CASE_									=						self.dropout(_A						,		_A)
      SCREAMING_SNAKE_CASE_									=						self.conva(_A)

      if self.conv_shortcut is not None:
          SCREAMING_SNAKE_CASE_									=						self.conv_shortcut(_A)

      return hidden_states + residual





import unittest

from transformers import load_tool

from .test_tools_common import ToolTesterMixin


UpperCamelCase__  :      int			=    "\nHugging Face was founded in 2016 by French entrepreneurs Clément Delangue, Julien Chaumond, and Thomas Wolf originally as a company that developed a chatbot app targeted at teenagers.[2] After open-sourcing the model behind the chatbot, the company pivoted to focus on being a platform for machine learning.\n\nIn March 2021, Hugging Face raised $40 million in a Series B funding round.[3]\n\nOn April 28, 2021, the company launched the BigScience Research Workshop in collaboration with several other research groups to release an open large language model.[4] In 2022, the workshop concluded with the announcement of BLOOM, a multilingual large language model with 176 billion parameters.[5]\n"



class       __snake_case  (    unittest.TestCase     , lowerCAmelCase__       ):

  def 						lowerCAmelCase__     (					self):
      SCREAMING_SNAKE_CASE_									=						load_tool('text-question-answering')
      self.tool.setup()
      SCREAMING_SNAKE_CASE_									=						load_tool('text-question-answering'						,		remote=_A)

  def 						lowerCAmelCase__     (					self):
      SCREAMING_SNAKE_CASE_									=						self.tool(_A						,		'What did Hugging Face do in April 2021?')
      self.assertEqual(_A						,		'launched the BigScience Research Workshop')

  def 						lowerCAmelCase__     (					self):
      SCREAMING_SNAKE_CASE_									=						self.remote_tool(_A						,		'What did Hugging Face do in April 2021?')
      self.assertEqual(_A						,		'launched the BigScience Research Workshop')

  def 						lowerCAmelCase__     (					self):
      SCREAMING_SNAKE_CASE_									=						self.tool(text=_A						,		question='What did Hugging Face do in April 2021?')
      self.assertEqual(_A						,		'launched the BigScience Research Workshop')


  def 						lowerCAmelCase__     (					self):
      SCREAMING_SNAKE_CASE_									=						self.remote_tool(text=_A						,		question='What did Hugging Face do in April 2021?')
      self.assertEqual(_A						,		'launched the BigScience Research Workshop')


