import unittest

from transformers import BertGenerationConfig, is_torch_available
from transformers.testing_utils import require_torch, slow, torch_device

from ...generation.test_utils import GenerationTesterMixin
from ...test_configuration_common import ConfigTester
from ...test_modeling_common import ModelTesterMixin, floats_tensor, ids_tensor, random_attention_mask
from ...test_pipeline_mixin import PipelineTesterMixin


if is_torch_available():
   import torch

   from transformers import BertGenerationDecoder, BertGenerationEncoder





class 		_lowercase						:

 def __init__( self	:							List[Any]   ,      snake_case	:							List[str]   ,      snake_case	:							Optional[int]=1_3   ,      snake_case	:							Any=7   ,      snake_case	:							List[Any]=True   ,      snake_case	:							Optional[Any]=True   ,      snake_case	:							str=9_9   ,      snake_case	:							Union[str, Any]=3_2   ,      snake_case	:							Dict=5   ,      snake_case	:							Dict=4   ,      snake_case	:							Optional[Any]=3_7   ,      snake_case	:							Optional[Any]="gelu"   ,      snake_case	:							Dict=0.1   ,      snake_case	:							Optional[int]=0.1   ,      snake_case	:							str=5_0   ,      snake_case	:							Optional[Any]=0.02   ,      snake_case	:							Optional[Any]=True   ,      snake_case	:							List[str]=None   ,      )   ->						List[str]:

  """simple docstring"""





  UpperCamelCase_					:						int =					parent
  UpperCamelCase_					:						Optional[int] =					batch_size
  UpperCamelCase_					:						List[str] =					seq_length
  UpperCamelCase_					:						Tuple =					is_training
  UpperCamelCase_					:						int =					use_input_mask
  UpperCamelCase_					:						int =					vocab_size
  UpperCamelCase_					:						Any =					hidden_size
  UpperCamelCase_					:						Dict =					num_hidden_layers
  UpperCamelCase_					:						str =					num_attention_heads
  UpperCamelCase_					:						Dict =					intermediate_size
  UpperCamelCase_					:						Tuple =					hidden_act
  UpperCamelCase_					:						Optional[Any] =					hidden_dropout_prob
  UpperCamelCase_					:						int =					attention_probs_dropout_prob
  UpperCamelCase_					:						str =					max_position_embeddings
  UpperCamelCase_					:						str =					initializer_range
  UpperCamelCase_					:						List[Any] =					use_labels
  UpperCamelCase_					:						Optional[Any] =					scope

 def 						SCREAMING_SNAKE_CASE__   ( self	:							List[Any]			)   ->						Optional[int]:

  """simple docstring"""





  UpperCamelCase_					:						Union[str, Any] =					ids_tensor([self.batch_size, self.seq_length]   ,      self.vocab_size			)

  UpperCamelCase_					:						Dict =					None
  if self.use_input_mask:
   UpperCamelCase_					:						List[Any] =					random_attention_mask([self.batch_size, self.seq_length]			)

  if self.use_labels:
   UpperCamelCase_					:						Any =					ids_tensor([self.batch_size, self.seq_length]   ,      self.vocab_size			)

  UpperCamelCase_					:						Any =					self.get_config()

  return config, input_ids, input_mask, token_labels

 def 						SCREAMING_SNAKE_CASE__   ( self	:							Dict			)   ->						Optional[Any]:

  """simple docstring"""





  return BertGenerationConfig(
      vocab_size=self.vocab_size   ,      hidden_size=self.hidden_size   ,      num_hidden_layers=self.num_hidden_layers   ,      num_attention_heads=self.num_attention_heads   ,      intermediate_size=self.intermediate_size   ,      hidden_act=self.hidden_act   ,      hidden_dropout_prob=self.hidden_dropout_prob   ,      attention_probs_dropout_prob=self.attention_probs_dropout_prob   ,      max_position_embeddings=self.max_position_embeddings   ,      is_decoder=snake_case   ,      initializer_range=self.initializer_range   ,      )

 def 						SCREAMING_SNAKE_CASE__   ( self	:							Optional[int]			)   ->						str:

  """simple docstring"""





  (
      (
      UpperCamelCase_
  ), (
      UpperCamelCase_
  ), (
      UpperCamelCase_
  ), (
      UpperCamelCase_
  ), 
  )					:						Optional[Any] =					self.prepare_config_and_inputs()

  UpperCamelCase_					:						List[str] =					True
  UpperCamelCase_					:						Union[str, Any] =					floats_tensor([self.batch_size, self.seq_length, self.hidden_size]			)
  UpperCamelCase_					:						List[str] =					ids_tensor([self.batch_size, self.seq_length]   ,      vocab_size=2			)

  return (
      config,
      input_ids,
      input_mask,
      token_labels,
      encoder_hidden_states,
      encoder_attention_mask,
  )

 def 						SCREAMING_SNAKE_CASE__   ( self	:							Any   ,      snake_case	:							Union[str, Any]   ,      snake_case	:							int   ,      snake_case	:							List[Any]   ,      snake_case	:							Any   ,      **snake_case	:							int   ,      )   ->						List[str]:

  """simple docstring"""





  UpperCamelCase_					:						Optional[Any] =					BertGenerationEncoder(config=snake_case			)
  model.to(snake_case			)
  model.eval()
  UpperCamelCase_					:						List[str] =					model(snake_case   ,      attention_mask=snake_case			)
  UpperCamelCase_					:						Dict =					model(snake_case			)
  self.parent.assertEqual(result.last_hidden_state.shape   ,      (self.batch_size, self.seq_length, self.hidden_size)			)

 def 						SCREAMING_SNAKE_CASE__   ( self	:							Any   ,      snake_case	:							Union[str, Any]   ,      snake_case	:							int   ,      snake_case	:							str   ,      snake_case	:							int   ,      snake_case	:							Tuple   ,      snake_case	:							Dict   ,      **snake_case	:							Union[str, Any]   ,      )   ->						Optional[int]:

  """simple docstring"""





  UpperCamelCase_					:						str =					True
  UpperCamelCase_					:						Optional[Any] =					BertGenerationEncoder(config=snake_case			)
  model.to(snake_case			)
  model.eval()
  UpperCamelCase_					:						str =					model(
      snake_case   ,      attention_mask=snake_case   ,      encoder_hidden_states=snake_case   ,      encoder_attention_mask=snake_case   ,      )
  UpperCamelCase_					:						Optional[int] =					model(
      snake_case   ,      attention_mask=snake_case   ,      encoder_hidden_states=snake_case   ,      )
  self.parent.assertEqual(result.last_hidden_state.shape   ,      (self.batch_size, self.seq_length, self.hidden_size)			)

 def 						SCREAMING_SNAKE_CASE__   ( self	:							int   ,      snake_case	:							Optional[Any]   ,      snake_case	:							Any   ,      snake_case	:							str   ,      snake_case	:							Union[str, Any]   ,      snake_case	:							List[str]   ,      snake_case	:							Optional[int]   ,      **snake_case	:							Dict   ,      )   ->						str:

  """simple docstring"""





  UpperCamelCase_					:						Tuple =					True
  UpperCamelCase_					:						Optional[int] =					True
  UpperCamelCase_					:						Optional[int] =					BertGenerationDecoder(config=snake_case			).to(snake_case			).eval()

  # first forward pass
  UpperCamelCase_					:						Optional[int] =					model(
      snake_case   ,      attention_mask=snake_case   ,      encoder_hidden_states=snake_case   ,      encoder_attention_mask=snake_case   ,      use_cache=snake_case   ,      )
  UpperCamelCase_					:						Union[str, Any] =					outputs.past_key_values

  # create hypothetical multiple next token and extent to next_input_ids
  UpperCamelCase_					:						Optional[Any] =					ids_tensor((self.batch_size, 3)   ,      config.vocab_size			)
  UpperCamelCase_					:						List[str] =					ids_tensor((self.batch_size, 3)   ,      vocab_size=2			)

  # append to next input_ids and
  UpperCamelCase_					:						Tuple =					torch.cat([input_ids, next_tokens]   ,      dim=-1			)
  UpperCamelCase_					:						List[Any] =					torch.cat([input_mask, next_mask]   ,      dim=-1			)

  UpperCamelCase_					:						Any =					model(
      snake_case   ,      attention_mask=snake_case   ,      encoder_hidden_states=snake_case   ,      encoder_attention_mask=snake_case   ,      output_hidden_states=snake_case   ,      )['hidden_states'][0]
  UpperCamelCase_					:						List[str] =					model(
      snake_case   ,      attention_mask=snake_case   ,      encoder_hidden_states=snake_case   ,      encoder_attention_mask=snake_case   ,      past_key_values=snake_case   ,      output_hidden_states=snake_case   ,      )['hidden_states'][0]

  # select random slice
  UpperCamelCase_					:						Optional[Any] =					ids_tensor((1,)   ,      output_from_past.shape[-1]			).item()
  UpperCamelCase_					:						Tuple =					output_from_no_past[:, -3:, random_slice_idx].detach()
  UpperCamelCase_					:						Optional[int] =					output_from_past[:, :, random_slice_idx].detach()

  self.parent.assertTrue(output_from_past_slice.shape[1] == next_tokens.shape[1]			)

  # test that outputs are equal for slice
  self.parent.assertTrue(torch.allclose(snake_case   ,      snake_case   ,      atol=1e-3			)			)

 def 						SCREAMING_SNAKE_CASE__   ( self	:							str   ,      snake_case	:							Dict   ,      snake_case	:							Dict   ,      snake_case	:							int   ,      snake_case	:							Optional[int]   ,      *snake_case	:							Dict   ,      )   ->						Any:

  """simple docstring"""





  UpperCamelCase_					:						Tuple =					BertGenerationDecoder(snake_case			)
  model.to(snake_case			)
  model.eval()
  UpperCamelCase_					:						str =					model(snake_case   ,      attention_mask=snake_case   ,      labels=snake_case			)
  self.parent.assertEqual(result.logits.shape   ,      (self.batch_size, self.seq_length, self.vocab_size)			)


 def 						SCREAMING_SNAKE_CASE__   ( self	:							Tuple			)   ->						Optional[int]:

  """simple docstring"""





  UpperCamelCase_, UpperCamelCase_, UpperCamelCase_, UpperCamelCase_					:						Any =					self.prepare_config_and_inputs()
  UpperCamelCase_					:						Dict =					{'input_ids': input_ids, 'attention_mask': input_mask}
  return config, inputs_dict





@require_torch
class 		_lowercase						(					snake_case_						,      snake_case_						,      snake_case_						,      unittest.TestCase       ):
 lowercase				     =   (BertGenerationEncoder, BertGenerationDecoder) if is_torch_available() else ()
 lowercase				     =   (BertGenerationDecoder,) if is_torch_available() else ()
 lowercase				     =   (
     {'feature-extraction': BertGenerationEncoder, 'text-generation': BertGenerationDecoder}
     if is_torch_available()
     else {}
 )

 def 						SCREAMING_SNAKE_CASE__   ( self	:							int			)   ->						List[Any]:

  """simple docstring"""





  UpperCamelCase_					:						Optional[Any] =					BertGenerationEncoderTester(self			)
  UpperCamelCase_					:						List[str] =					ConfigTester(self   ,      config_class=snake_case   ,      hidden_size=3_7			)

 def 						SCREAMING_SNAKE_CASE__   ( self	:							Any			)   ->						str:

  """simple docstring"""





  self.config_tester.run_common_tests()

 def 						SCREAMING_SNAKE_CASE__   ( self	:							str			)   ->						Optional[Any]:

  """simple docstring"""





  UpperCamelCase_					:						int =					self.model_tester.prepare_config_and_inputs()
  self.model_tester.create_and_check_model(*snake_case			)

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

  """simple docstring"""





  UpperCamelCase_, UpperCamelCase_, UpperCamelCase_, UpperCamelCase_					:						Optional[int] =					self.model_tester.prepare_config_and_inputs()
  UpperCamelCase_					:						Optional[int] =					'bert'
  self.model_tester.create_and_check_model(snake_case   ,      snake_case   ,      snake_case   ,      snake_case			)

 def 						SCREAMING_SNAKE_CASE__   ( self	:							List[Any]			)   ->						Tuple:

  """simple docstring"""





  UpperCamelCase_					:						Optional[Any] =					self.model_tester.prepare_config_and_inputs_for_decoder()
  self.model_tester.create_and_check_model_as_decoder(*snake_case			)

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

  """simple docstring"""





  UpperCamelCase_					:						List[Any] =					self.model_tester.prepare_config_and_inputs_for_decoder()
  self.model_tester.create_and_check_decoder_model_past_large_inputs(*snake_case			)

 def 						SCREAMING_SNAKE_CASE__   ( self	:							str			)   ->						Optional[Any]:

  """simple docstring"""





  (
      (
      UpperCamelCase_
  ), (
      UpperCamelCase_
  ), (
      UpperCamelCase_
  ), (
      UpperCamelCase_
  ), (
      UpperCamelCase_
  ), (
      UpperCamelCase_
  ), 
  )					:						Optional[int] =					self.model_tester.prepare_config_and_inputs_for_decoder()

  UpperCamelCase_					:						Optional[Any] =					None

  self.model_tester.create_and_check_model_as_decoder(
      snake_case   ,      snake_case   ,      snake_case   ,      snake_case   ,      snake_case   ,      snake_case   ,      )

 def 						SCREAMING_SNAKE_CASE__   ( self	:							Optional[Any]			)   ->						Optional[Any]:

  """simple docstring"""





  UpperCamelCase_					:						int =					self.model_tester.prepare_config_and_inputs_for_decoder()
  self.model_tester.create_and_check_for_causal_lm(*snake_case			)


 @slow
 def 						SCREAMING_SNAKE_CASE__   ( self	:							Optional[Any]			)   ->						Optional[int]:

  """simple docstring"""





  UpperCamelCase_					:						List[Any] =					BertGenerationEncoder.from_pretrained('google/bert_for_seq_generation_L-24_bbc_encoder'			)
  self.assertIsNotNone(snake_case			)





@require_torch
class 		_lowercase						(					unittest.TestCase       ):

 @slow
 def 						SCREAMING_SNAKE_CASE__   ( self	:							int			)   ->						Any:

  """simple docstring"""





  UpperCamelCase_					:						Optional[int] =					BertGenerationEncoder.from_pretrained('google/bert_for_seq_generation_L-24_bbc_encoder'			)
  UpperCamelCase_					:						Optional[int] =					torch.tensor([[1_0_1, 7_5_9_2, 1_0_1_0, 2_0_2_6, 3_8_9_9, 2_0_0_3, 1_0_1_4_0, 1_0_2]]			)
  with torch.no_grad():
   UpperCamelCase_					:						str =					model(snake_case			)[0]
  UpperCamelCase_					:						Optional[int] =					torch.Size([1, 8, 1_0_2_4]			)
  self.assertEqual(output.shape   ,      snake_case			)
  UpperCamelCase_					:						Dict =					torch.tensor(
      [[[0.1775, 0.0083, -0.0321], [1.6002, 0.1287, 0.3912], [2.1473, 0.5791, 0.6066]]]			)
  self.assertTrue(torch.allclose(output[:, :3, :3]   ,      snake_case   ,      atol=1e-4			)			)





@require_torch
class 		_lowercase						(					unittest.TestCase       ):

 @slow
 def 						SCREAMING_SNAKE_CASE__   ( self	:							str			)   ->						Union[str, Any]:

  """simple docstring"""





  UpperCamelCase_					:						int =					BertGenerationDecoder.from_pretrained('google/bert_for_seq_generation_L-24_bbc_encoder'			)
  UpperCamelCase_					:						Any =					torch.tensor([[1_0_1, 7_5_9_2, 1_0_1_0, 2_0_2_6, 3_8_9_9, 2_0_0_3, 1_0_1_4_0, 1_0_2]]			)
  with torch.no_grad():
   UpperCamelCase_					:						Dict =					model(snake_case			)[0]
  UpperCamelCase_					:						Optional[int] =					torch.Size([1, 8, 5_0_3_5_8]			)
  self.assertEqual(output.shape   ,      snake_case			)
  UpperCamelCase_					:						int =					torch.tensor(
      [[[-0.5788, -2.5994, -3.7054], [0.0438, 4.7997, 1.8795], [1.5862, 6.6409, 4.4638]]]			)
  self.assertTrue(torch.allclose(output[:, :3, :3]   ,      snake_case   ,      atol=1e-4			)			)




import gc
import unittest

from diffusers import FlaxDPMSolverMultistepScheduler, FlaxStableDiffusionPipeline
from diffusers.utils import is_flax_available, slow
from diffusers.utils.testing_utils import require_flax


if is_flax_available():
   import jax
   import jax.numpy as jnp
   from flax.jax_utils import replicate
   from flax.training.common_utils import shard





@slow
@require_flax
class 		_lowercase						(					unittest.TestCase       ):

 def 						SCREAMING_SNAKE_CASE__   ( self	:							int			)   ->						str:

  """simple docstring"""





  super().tearDown()
  gc.collect()

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

  """simple docstring"""





  UpperCamelCase_, UpperCamelCase_					:						str =					FlaxStableDiffusionPipeline.from_pretrained(
      'stabilityai/stable-diffusion-2'   ,      revision='bf16'   ,      dtype=jnp.bfloataa   ,      )

  UpperCamelCase_					:						str =					'A painting of a squirrel eating a burger'
  UpperCamelCase_					:						Any =					jax.device_count()
  UpperCamelCase_					:						List[str] =					num_samples * [prompt]
  UpperCamelCase_					:						List[Any] =					sd_pipe.prepare_inputs(snake_case			)

  UpperCamelCase_					:						Dict =					replicate(snake_case			)
  UpperCamelCase_					:						Optional[Any] =					shard(snake_case			)

  UpperCamelCase_					:						Dict =					jax.random.PRNGKey(0			)
  UpperCamelCase_					:						Tuple =					jax.random.split(snake_case   ,      jax.device_count()			)

  UpperCamelCase_					:						Optional[Any] =					sd_pipe(snake_case   ,      snake_case   ,      snake_case   ,      num_inference_steps=2_5   ,      jit=snake_case			)[0]
  assert images.shape == (jax.device_count(), 1, 7_6_8, 7_6_8, 3)

  UpperCamelCase_					:						int =					images.reshape((images.shape[0] * images.shape[1],) + images.shape[-3:]			)
  UpperCamelCase_					:						Tuple =					images[0, 2_5_3:2_5_6, 2_5_3:2_5_6, -1]

  UpperCamelCase_					:						str =					jnp.asarray(jax.device_get(image_slice.flatten()			)			)
  UpperCamelCase_					:						List[str] =					jnp.array([0.4238, 0.4414, 0.4395, 0.4453, 0.4629, 0.4590, 0.4531, 0.45508, 0.4512]			)
  print(f"output_slice: {output_slice}"			)
  assert jnp.abs(output_slice - expected_slice			).max() < 1e-2


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

  """simple docstring"""





  UpperCamelCase_					:						Tuple =					'stabilityai/stable-diffusion-2'
  UpperCamelCase_, UpperCamelCase_					:						Tuple =					FlaxDPMSolverMultistepScheduler.from_pretrained(snake_case   ,      subfolder='scheduler'			)
  UpperCamelCase_, UpperCamelCase_					:						Union[str, Any] =					FlaxStableDiffusionPipeline.from_pretrained(
      snake_case   ,      scheduler=snake_case   ,      revision='bf16'   ,      dtype=jnp.bfloataa   ,      )
  UpperCamelCase_					:						Optional[int] =					scheduler_params

  UpperCamelCase_					:						Optional[Any] =					'A painting of a squirrel eating a burger'
  UpperCamelCase_					:						Union[str, Any] =					jax.device_count()
  UpperCamelCase_					:						Union[str, Any] =					num_samples * [prompt]
  UpperCamelCase_					:						Tuple =					sd_pipe.prepare_inputs(snake_case			)

  UpperCamelCase_					:						List[Any] =					replicate(snake_case			)
  UpperCamelCase_					:						Optional[Any] =					shard(snake_case			)

  UpperCamelCase_					:						Tuple =					jax.random.PRNGKey(0			)
  UpperCamelCase_					:						str =					jax.random.split(snake_case   ,      jax.device_count()			)

  UpperCamelCase_					:						str =					sd_pipe(snake_case   ,      snake_case   ,      snake_case   ,      num_inference_steps=2_5   ,      jit=snake_case			)[0]
  assert images.shape == (jax.device_count(), 1, 7_6_8, 7_6_8, 3)

  UpperCamelCase_					:						str =					images.reshape((images.shape[0] * images.shape[1],) + images.shape[-3:]			)
  UpperCamelCase_					:						int =					images[0, 2_5_3:2_5_6, 2_5_3:2_5_6, -1]

  UpperCamelCase_					:						int =					jnp.asarray(jax.device_get(image_slice.flatten()			)			)
  UpperCamelCase_					:						Union[str, Any] =					jnp.array([0.4336, 0.42969, 0.4453, 0.4199, 0.4297, 0.4531, 0.4434, 0.4434, 0.4297]			)
  print(f"output_slice: {output_slice}"			)
  assert jnp.abs(output_slice - expected_slice			).max() < 1e-2



