
import unittest

import numpy as np

from transformers import AlbertConfig, is_flax_available
from transformers.testing_utils import require_flax, slow

from ...test_modeling_flax_common import FlaxModelTesterMixin, ids_tensor, random_attention_mask


if is_flax_available():
      import jax.numpy as jnp

      from transformers.models.albert.modeling_flax_albert import (
          FlaxAlbertForMaskedLM,
          FlaxAlbertForMultipleChoice,
          FlaxAlbertForPreTraining,
          FlaxAlbertForQuestionAnswering,
          FlaxAlbertForSequenceClassification,
          FlaxAlbertForTokenClassification,
          FlaxAlbertModel,
      )



class 		UpperCamelCase__  (       unittest.TestCase    ):
   def __init__(self :					Dict					,				snake_case_ :					str					,				snake_case_ :					Union[str, Any]=1_3					,				snake_case_ :					Optional[int]=7					,				snake_case_ :					Tuple=True					,				snake_case_ :					int=True					,				snake_case_ :					Dict=True					,				snake_case_ :					int=True					,				snake_case_ :					int=9_9					,				snake_case_ :					Optional[Any]=3_2					,				snake_case_ :					Optional[int]=5					,				snake_case_ :					Optional[int]=4					,				snake_case_ :					Optional[int]=3_7					,				snake_case_ :					Tuple="gelu"					,				snake_case_ :					Dict=0.1					,				snake_case_ :					Dict=0.1					,				snake_case_ :					int=5_1_2					,				snake_case_ :					int=1_6					,				snake_case_ :					int=2					,				snake_case_ :					str=0.02					,				snake_case_ :					Optional[int]=4					,				):
    __a :  Any						=  parent
    __a :  str						=  batch_size
    __a :  List[str]						=  seq_length
    __a :  Dict						=  is_training
    __a :  int						=  use_attention_mask
    __a :  Tuple						=  use_token_type_ids
    __a :  int						=  use_labels
    __a :  List[str]						=  vocab_size
    __a :  Tuple						=  hidden_size
    __a :  int						=  num_hidden_layers
    __a :  Dict						=  num_attention_heads
    __a :  List[str]						=  intermediate_size
    __a :  str						=  hidden_act
    __a :  int						=  hidden_dropout_prob
    __a :  List[str]						=  attention_probs_dropout_prob
    __a :  List[str]						=  max_position_embeddings
    __a :  Union[str, Any]						=  type_vocab_size
    __a :  Any						=  type_sequence_label_size
    __a :  Dict						=  initializer_range
    __a :  str						=  num_choices
   def       lowerCAmelCase			(self :					Optional[int]   ):
    __a :  Optional[int]						=  ids_tensor([self.batch_size, self.seq_length]					,				self.vocab_size   )

    __a :  Optional[Any]						=  None
    if self.use_attention_mask:
     __a :  Tuple						=  random_attention_mask([self.batch_size, self.seq_length]   )

    __a :  Dict						=  None
    if self.use_token_type_ids:
     __a :  str						=  ids_tensor([self.batch_size, self.seq_length]					,				self.type_vocab_size   )

    __a :  Optional[Any]						=  AlbertConfig(
        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					,				type_vocab_size=self.type_vocab_size					,				is_decoder=snake_case_					,				initializer_range=self.initializer_range					,				)

    return config, input_ids, token_type_ids, attention_mask







   def       lowerCAmelCase			(self :					str   ):
    __a :  List[Any]						=  self.prepare_config_and_inputs()
    __a				,       __a				,       __a				,       __a :  Union[str, Any]						=  config_and_inputs
    __a :  List[str]						=  {'''input_ids''': input_ids, '''token_type_ids''': token_type_ids, '''attention_mask''': attention_mask}
    return config, inputs_dict





@require_flax
class 		UpperCamelCase__  (       __lowercase    ,unittest.TestCase    ):
   _SCREAMING_SNAKE_CASE      :							Dict	       =      (
       (
           FlaxAlbertModel,
           FlaxAlbertForPreTraining,
           FlaxAlbertForMaskedLM,
           FlaxAlbertForMultipleChoice,
           FlaxAlbertForQuestionAnswering,
           FlaxAlbertForSequenceClassification,
           FlaxAlbertForTokenClassification,
           FlaxAlbertForQuestionAnswering,
       )
       if is_flax_available()
       else ()
   )
   def       lowerCAmelCase			(self :					Union[str, Any]   ):
    __a :  Union[str, Any]						=  FlaxAlbertModelTester(self   )







   @slow
   def       lowerCAmelCase			(self :					Optional[Any]   ):
    for model_class_name in self.all_model_classes:
     __a :  Tuple						=  model_class_name.from_pretrained('''albert-base-v2'''   )
     __a :  Any						=  model(np.ones((1, 1)   )   )
     self.assertIsNotNone(snake_case_   )





@require_flax
class 		UpperCamelCase__  (       unittest.TestCase    ):
   @slow
   def       lowerCAmelCase			(self :					Optional[int]   ):
    __a :  Dict						=  FlaxAlbertModel.from_pretrained('''albert-base-v2'''   )
    __a :  Tuple						=  np.array([[0, 3_4_5, 2_3_2, 3_2_8, 7_4_0, 1_4_0, 1_6_9_5, 6_9, 6_0_7_8, 1_5_8_8, 2]]   )
    __a :  Any						=  np.array([[0, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1]]   )
    __a :  List[str]						=  model(snake_case_					,				attention_mask=snake_case_   )[0]
    __a :  Any						=  (1, 1_1, 7_6_8)
    self.assertEqual(output.shape					,				snake_case_   )
    __a :  str						=  np.array(
        [[[-0.6513, 1.5035, -0.2766], [-0.6515, 1.5046, -0.2780], [-0.6512, 1.5049, -0.2784]]]   )

    self.assertTrue(jnp.allclose(output[:, 1:4, 1:4]					,				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 		UpperCamelCase__  (       unittest.TestCase    ):
   def       lowerCAmelCase			(self :					Optional[Any]   ):
    # clean up the VRAM after each test
    super().tearDown()
    gc.collect()
   def       lowerCAmelCase			(self :					int   ):
    __a				,       __a :  Optional[Any]						=  FlaxStableDiffusionPipeline.from_pretrained(
        '''stabilityai/stable-diffusion-2'''					,				revision='''bf16'''					,				dtype=jnp.bfloataa					,				)

    __a :  Any						=  '''A painting of a squirrel eating a burger'''
    __a :  Dict						=  jax.device_count()
    __a :  Optional[int]						=  num_samples * [prompt]
    __a :  Optional[Any]						=  sd_pipe.prepare_inputs(snake_case_   )

    __a :  Optional[Any]						=  replicate(snake_case_   )
    __a :  Optional[int]						=  shard(snake_case_   )

    __a :  int						=  jax.random.PRNGKey(0   )
    __a :  str						=  jax.random.split(snake_case_					,				jax.device_count()   )

    __a :  int						=  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)

    __a :  Optional[int]						=  images.reshape((images.shape[0] * images.shape[1],) + images.shape[-3:]   )
    __a :  List[Any]						=  images[0, 2_5_3:2_5_6, 2_5_3:2_5_6, -1]

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







   def       lowerCAmelCase			(self :					Tuple   ):
    __a :  Optional[Any]						=  '''stabilityai/stable-diffusion-2'''
    __a				,       __a :  Dict						=  FlaxDPMSolverMultistepScheduler.from_pretrained(snake_case_					,				subfolder='''scheduler'''   )
    __a				,       __a :  List[Any]						=  FlaxStableDiffusionPipeline.from_pretrained(
        snake_case_					,				scheduler=snake_case_					,				revision='''bf16'''					,				dtype=jnp.bfloataa					,				)
    __a :  Union[str, Any]						=  scheduler_params

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

    __a :  Tuple						=  replicate(snake_case_   )
    __a :  Dict						=  shard(snake_case_   )

    __a :  Dict						=  jax.random.PRNGKey(0   )
    __a :  Dict						=  jax.random.split(snake_case_					,				jax.device_count()   )

    __a :  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)

    __a :  List[Any]						=  images.reshape((images.shape[0] * images.shape[1],) + images.shape[-3:]   )
    __a :  Any						=  images[0, 2_5_3:2_5_6, 2_5_3:2_5_6, -1]

    __a :  List[Any]						=  jnp.asarray(jax.device_get(image_slice.flatten()   )   )
    __a :  List[Any]						=  jnp.array([0.4336, 0.4_2969, 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
