





"""simple docstring"""



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 		__snake_case					:





 """simple docstring"""


 def __init__(    self					,						__lowerCamelCase					,						__lowerCamelCase=13					,						__lowerCamelCase=7					,						__lowerCamelCase=True					,						__lowerCamelCase=True					,						__lowerCamelCase=99					,						__lowerCamelCase=32					,						__lowerCamelCase=5					,						__lowerCamelCase=4					,						__lowerCamelCase=37					,						__lowerCamelCase="gelu"					,						__lowerCamelCase=0.1					,						__lowerCamelCase=0.1					,						__lowerCamelCase=50					,						__lowerCamelCase=0.0_2					,						__lowerCamelCase=True					,						__lowerCamelCase=None					,						):



     '''simple docstring'''




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

 def 					UpperCamelCase__(    self  ):



     '''simple docstring'''




     __A    :						Optional[Any]						 =     ids_tensor([self.batch_size, self.seq_length]					,						self.vocab_size  )

     __A    :						Union[str, Any]						 =     None
     if self.use_input_mask:
         __A    :						Dict						 =     random_attention_mask([self.batch_size, self.seq_length]  )

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

     __A    :						Optional[Any]						 =     self.get_config()

     return config, input_ids, input_mask, token_labels

 def 					UpperCamelCase__(    self  ):



     '''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=__lowerCamelCase					,						initializer_range=self.initializer_range					,						)

 def 					UpperCamelCase__(    self  ):



     '''simple docstring'''




     (
         (
         __A
     ) ,       (
         __A
     ) ,       (
         __A
     ) ,       (
         __A
     ) ,       
     )    :						List[Any]						 =     self.prepare_config_and_inputs()

     __A    :						Tuple						 =     True
     __A    :						Dict						 =     floats_tensor([self.batch_size, self.seq_length, self.hidden_size]  )
     __A    :						Dict						 =     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 					UpperCamelCase__(    self					,						__lowerCamelCase					,						__lowerCamelCase					,						__lowerCamelCase					,						__lowerCamelCase					,						**__lowerCamelCase					,						):



     '''simple docstring'''




     __A    :						str						 =     BertGenerationEncoder(config=__lowerCamelCase  )
     model.to(__lowerCamelCase  )
     model.eval()
     __A    :						Union[str, Any]						 =     model(__lowerCamelCase					,						attention_mask=__lowerCamelCase  )
     __A    :						Union[str, Any]						 =     model(__lowerCamelCase  )
     self.parent.assertEqual(result.last_hidden_state.shape					,						(self.batch_size, self.seq_length, self.hidden_size)  )

 def 					UpperCamelCase__(    self					,						__lowerCamelCase					,						__lowerCamelCase					,						__lowerCamelCase					,						__lowerCamelCase					,						__lowerCamelCase					,						__lowerCamelCase					,						**__lowerCamelCase					,						):



     '''simple docstring'''




     __A    :						Tuple						 =     True
     __A    :						List[str]						 =     BertGenerationEncoder(config=__lowerCamelCase  )
     model.to(__lowerCamelCase  )
     model.eval()
     __A    :						Tuple						 =     model(
         __lowerCamelCase					,						attention_mask=__lowerCamelCase					,						encoder_hidden_states=__lowerCamelCase					,						encoder_attention_mask=__lowerCamelCase					,						)
     __A    :						int						 =     model(
         __lowerCamelCase					,						attention_mask=__lowerCamelCase					,						encoder_hidden_states=__lowerCamelCase					,						)
     self.parent.assertEqual(result.last_hidden_state.shape					,						(self.batch_size, self.seq_length, self.hidden_size)  )

 def 					UpperCamelCase__(    self					,						__lowerCamelCase					,						__lowerCamelCase					,						__lowerCamelCase					,						__lowerCamelCase					,						__lowerCamelCase					,						__lowerCamelCase					,						**__lowerCamelCase					,						):



     '''simple docstring'''




     __A    :						List[str]						 =     True
     __A    :						List[str]						 =     True
     __A    :						Optional[int]						 =     BertGenerationDecoder(config=__lowerCamelCase  ).to(__lowerCamelCase  ).eval()

     # first forward pass
     __A    :						Optional[Any]						 =     model(
         __lowerCamelCase					,						attention_mask=__lowerCamelCase					,						encoder_hidden_states=__lowerCamelCase					,						encoder_attention_mask=__lowerCamelCase					,						use_cache=__lowerCamelCase					,						)
     __A    :						int						 =     outputs.past_key_values

     # create hypothetical multiple next token and extent to next_input_ids
     __A    :						Tuple						 =     ids_tensor((self.batch_size, 3)					,						config.vocab_size  )
     __A    :						Dict						 =     ids_tensor((self.batch_size, 3)					,						vocab_size=2  )

     # append to next input_ids and
     __A    :						int						 =     torch.cat([input_ids, next_tokens]					,						dim=-1  )
     __A    :						Union[str, Any]						 =     torch.cat([input_mask, next_mask]					,						dim=-1  )

     __A    :						List[str]						 =     model(
         __lowerCamelCase					,						attention_mask=__lowerCamelCase					,						encoder_hidden_states=__lowerCamelCase					,						encoder_attention_mask=__lowerCamelCase					,						output_hidden_states=__lowerCamelCase					,						)['''hidden_states'''][0]
     __A    :						Tuple						 =     model(
         __lowerCamelCase					,						attention_mask=__lowerCamelCase					,						encoder_hidden_states=__lowerCamelCase					,						encoder_attention_mask=__lowerCamelCase					,						past_key_values=__lowerCamelCase					,						output_hidden_states=__lowerCamelCase					,						)['''hidden_states'''][0]

     # select random slice
     __A    :						str						 =     ids_tensor((1,)					,						output_from_past.shape[-1]  ).item()
     __A    :						List[str]						 =     output_from_no_past[:, -3:, random_slice_idx].detach()
     __A    :						List[str]						 =     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(__lowerCamelCase					,						__lowerCamelCase					,						atol=1e-3  )  )

 def 					UpperCamelCase__(    self					,						__lowerCamelCase					,						__lowerCamelCase					,						__lowerCamelCase					,						__lowerCamelCase					,						*__lowerCamelCase					,						):



     '''simple docstring'''




     __A    :						Optional[Any]						 =     BertGenerationDecoder(__lowerCamelCase  )
     model.to(__lowerCamelCase  )
     model.eval()
     __A    :						Optional[Any]						 =     model(__lowerCamelCase					,						attention_mask=__lowerCamelCase					,						labels=__lowerCamelCase  )
     self.parent.assertEqual(result.logits.shape					,						(self.batch_size, self.seq_length, self.vocab_size)  )







 def 					UpperCamelCase__(    self  ):



     '''simple docstring'''




     __A ,       __A ,       __A ,       __A    :						int						 =     self.prepare_config_and_inputs()
     __A    :						List[str]						 =     {'''input_ids''': input_ids, '''attention_mask''': input_mask}
     return config, inputs_dict






@require_torch
class 		__snake_case					(		SCREAMING_SNAKE_CASE__							,							SCREAMING_SNAKE_CASE__							,							SCREAMING_SNAKE_CASE__							,							unittest.TestCase     ):





 """simple docstring"""


 _lowerCamelCase        =     (BertGenerationEncoder, BertGenerationDecoder) if is_torch_available() else ()
 _lowerCamelCase        =     (BertGenerationDecoder,) if is_torch_available() else ()
 _lowerCamelCase        =     (
     {"""feature-extraction""": BertGenerationEncoder, """text-generation""": BertGenerationDecoder}
     if is_torch_available()
     else {}
 )

 def 					UpperCamelCase__(    self  ):



     '''simple docstring'''




     __A    :						Optional[Any]						 =     BertGenerationEncoderTester(self  )
     __A    :						Union[str, Any]						 =     ConfigTester(self					,						config_class=__lowerCamelCase					,						hidden_size=37  )

 def 					UpperCamelCase__(    self  ):



     '''simple docstring'''




     self.config_tester.run_common_tests()

 def 					UpperCamelCase__(    self  ):



     '''simple docstring'''




     __A    :						Tuple						 =     self.model_tester.prepare_config_and_inputs()
     self.model_tester.create_and_check_model(*__lowerCamelCase  )

 def 					UpperCamelCase__(    self  ):



     '''simple docstring'''




     __A ,       __A ,       __A ,       __A    :						Tuple						 =     self.model_tester.prepare_config_and_inputs()
     __A    :						Optional[Any]						 =     '''bert'''
     self.model_tester.create_and_check_model(__lowerCamelCase					,						__lowerCamelCase					,						__lowerCamelCase					,						__lowerCamelCase  )

 def 					UpperCamelCase__(    self  ):



     '''simple docstring'''




     __A    :						str						 =     self.model_tester.prepare_config_and_inputs_for_decoder()
     self.model_tester.create_and_check_model_as_decoder(*__lowerCamelCase  )

 def 					UpperCamelCase__(    self  ):



     '''simple docstring'''




     __A    :						Any						 =     self.model_tester.prepare_config_and_inputs_for_decoder()
     self.model_tester.create_and_check_decoder_model_past_large_inputs(*__lowerCamelCase  )

 def 					UpperCamelCase__(    self  ):



     '''simple docstring'''




     (
         (
         __A
     ) ,       (
         __A
     ) ,       (
         __A
     ) ,       (
         __A
     ) ,       (
         __A
     ) ,       (
         __A
     ) ,       
     )    :						List[str]						 =     self.model_tester.prepare_config_and_inputs_for_decoder()

     __A    :						Dict						 =     None

     self.model_tester.create_and_check_model_as_decoder(
         __lowerCamelCase					,						__lowerCamelCase					,						__lowerCamelCase					,						__lowerCamelCase					,						__lowerCamelCase					,						__lowerCamelCase					,						)

 def 					UpperCamelCase__(    self  ):



     '''simple docstring'''




     __A    :						Union[str, Any]						 =     self.model_tester.prepare_config_and_inputs_for_decoder()
     self.model_tester.create_and_check_for_causal_lm(*__lowerCamelCase  )







 @slow
 def 					UpperCamelCase__(    self  ):



     '''simple docstring'''




     __A    :						List[str]						 =     BertGenerationEncoder.from_pretrained('''google/bert_for_seq_generation_L-24_bbc_encoder'''  )
     self.assertIsNotNone(__lowerCamelCase  )






@require_torch
class 		__snake_case					(		unittest.TestCase     ):





 """simple docstring"""


 @slow
 def 					UpperCamelCase__(    self  ):



     '''simple docstring'''




     __A    :						Optional[Any]						 =     BertGenerationEncoder.from_pretrained('''google/bert_for_seq_generation_L-24_bbc_encoder'''  )
     __A    :						int						 =     torch.tensor([[101, 7592, 1010, 2026, 3899, 2003, 1_0140, 102]]  )
     with torch.no_grad():
         __A    :						Optional[Any]						 =     model(__lowerCamelCase  )[0]
     __A    :						Optional[Any]						 =     torch.Size([1, 8, 1024]  )
     self.assertEqual(output.shape					,						__lowerCamelCase  )
     __A    :						Dict						 =     torch.tensor(
         [[[0.1_7_7_5, 0.0_0_8_3, -0.0_3_2_1], [1.6_0_0_2, 0.1_2_8_7, 0.3_9_1_2], [2.1_4_7_3, 0.5_7_9_1, 0.6_0_6_6]]]  )
     self.assertTrue(torch.allclose(output[:, :3, :3]					,						__lowerCamelCase					,						atol=1e-4  )  )






@require_torch
class 		__snake_case					(		unittest.TestCase     ):





 """simple docstring"""


 @slow
 def 					UpperCamelCase__(    self  ):



     '''simple docstring'''




     __A    :						Union[str, Any]						 =     BertGenerationDecoder.from_pretrained('''google/bert_for_seq_generation_L-24_bbc_encoder'''  )
     __A    :						Union[str, Any]						 =     torch.tensor([[101, 7592, 1010, 2026, 3899, 2003, 1_0140, 102]]  )
     with torch.no_grad():
         __A    :						Dict						 =     model(__lowerCamelCase  )[0]
     __A    :						Any						 =     torch.Size([1, 8, 5_0358]  )
     self.assertEqual(output.shape					,						__lowerCamelCase  )
     __A    :						int						 =     torch.tensor(
         [[[-0.5_7_8_8, -2.5_9_9_4, -3.7_0_5_4], [0.0_4_3_8, 4.7_9_9_7, 1.8_7_9_5], [1.5_8_6_2, 6.6_4_0_9, 4.4_6_3_8]]]  )
     self.assertTrue(torch.allclose(output[:, :3, :3]					,						__lowerCamelCase					,						atol=1e-4  )  )











"""simple docstring"""



# DISCLAIMER: This file is strongly influenced by https://github.com/ermongroup/ddim

from dataclasses import dataclass
from typing import Optional, Tuple, Union

import flax
import jax
import jax.numpy as jnp

from ..configuration_utils import ConfigMixin, register_to_config
from .scheduling_utils_flax import (
    CommonSchedulerState,
    FlaxKarrasDiffusionSchedulers,
    FlaxSchedulerMixin,
    FlaxSchedulerOutput,
    add_noise_common,
    get_velocity_common,
)






@flax.struct.dataclass
class 		__snake_case					:





 """simple docstring"""


 _lowerCamelCase        =     42

 # setable values
 _lowerCamelCase        =     42
 _lowerCamelCase        =     42
 _lowerCamelCase        =     None

 @classmethod
 def 					UpperCamelCase__(    cls					,						__lowerCamelCase					,						__lowerCamelCase					,						__lowerCamelCase  ):



     '''simple docstring'''




     return cls(common=__lowerCamelCase					,						init_noise_sigma=__lowerCamelCase					,						timesteps=__lowerCamelCase  )






@dataclass
class 		__snake_case					(		SCREAMING_SNAKE_CASE__     ):





 """simple docstring"""


 _lowerCamelCase        =     42






class 		__snake_case					(		SCREAMING_SNAKE_CASE__							,							SCREAMING_SNAKE_CASE__     ):





 """simple docstring"""


 _lowerCamelCase        =     [e.name for e in FlaxKarrasDiffusionSchedulers]

 _lowerCamelCase        =     42

 @property
 def 					UpperCamelCase__(    self  ):



     '''simple docstring'''




     return True

 @register_to_config
 def __init__(    self					,						__lowerCamelCase = 1000					,						__lowerCamelCase = 0.0_0_0_1					,						__lowerCamelCase = 0.0_2					,						__lowerCamelCase = "linear"					,						__lowerCamelCase = None					,						__lowerCamelCase = "fixed_small"					,						__lowerCamelCase = True					,						__lowerCamelCase = "epsilon"					,						__lowerCamelCase = jnp.floataa					,						):



     '''simple docstring'''




     __A    :						Tuple						 =     dtype

 def 					UpperCamelCase__(    self					,						__lowerCamelCase = None  ):



     '''simple docstring'''




     if common is None:
         __A    :						Tuple						 =     CommonSchedulerState.create(self  )

     # standard deviation of the initial noise distribution
     __A    :						Tuple						 =     jnp.array(1.0					,						dtype=self.dtype  )

     __A    :						Optional[int]						 =     jnp.arange(0					,						self.config.num_train_timesteps  ).round()[::-1]

     return DDPMSchedulerState.create(
         common=__lowerCamelCase					,						init_noise_sigma=__lowerCamelCase					,						timesteps=__lowerCamelCase					,						)

 def 					UpperCamelCase__(    self					,						__lowerCamelCase					,						__lowerCamelCase					,						__lowerCamelCase = None  ):



     '''simple docstring'''




     return sample

 def 					UpperCamelCase__(    self					,						__lowerCamelCase					,						__lowerCamelCase					,						__lowerCamelCase = ()  ):



     '''simple docstring'''




     __A    :						Optional[Any]						 =     self.config.num_train_timesteps // num_inference_steps
     # creates integer timesteps by multiplying by ratio
     # rounding to avoid issues when num_inference_step is power of 3
     __A    :						Optional[Any]						 =     (jnp.arange(0					,						__lowerCamelCase  ) * step_ratio).round()[::-1]

     return state.replace(
         num_inference_steps=__lowerCamelCase					,						timesteps=__lowerCamelCase					,						)

 def 					UpperCamelCase__(    self					,						__lowerCamelCase					,						__lowerCamelCase					,						__lowerCamelCase=None					,						__lowerCamelCase=None  ):



     '''simple docstring'''




     __A    :						int						 =     state.common.alphas_cumprod[t]
     __A    :						List[str]						 =     jnp.where(t > 0					,						state.common.alphas_cumprod[t - 1]					,						jnp.array(1.0					,						dtype=self.dtype  )  )

     # For t > 0, compute predicted variance βt (see formula (6) and (7) from https://arxiv.org/pdf/2006.11239.pdf)
     # and sample from it to get previous sample
     # x_{t-1} ~ N(pred_prev_sample, variance) == add variance to pred_sample
     __A    :						str						 =     (1 - alpha_prod_t_prev) / (1 - alpha_prod_t) * state.common.betas[t]

     if variance_type is None:
         __A    :						Dict						 =     self.config.variance_type

     # hacks - were probably added for training stability
     if variance_type == "fixed_small":
         __A    :						List[Any]						 =     jnp.clip(__lowerCamelCase					,						a_min=1e-2_0  )
     # for rl-diffuser https://arxiv.org/abs/2205.09991
     elif variance_type == "fixed_small_log":
         __A    :						Optional[Any]						 =     jnp.log(jnp.clip(__lowerCamelCase					,						a_min=1e-2_0  )  )
     elif variance_type == "fixed_large":
         __A    :						Tuple						 =     state.common.betas[t]
     elif variance_type == "fixed_large_log":
         # Glide max_log
         __A    :						Union[str, Any]						 =     jnp.log(state.common.betas[t]  )
     elif variance_type == "learned":
         return predicted_variance
     elif variance_type == "learned_range":
         __A    :						Optional[Any]						 =     variance
         __A    :						Optional[Any]						 =     state.common.betas[t]
         __A    :						Any						 =     (predicted_variance + 1) / 2
         __A    :						Union[str, Any]						 =     frac * max_log + (1 - frac) * min_log

     return variance

 def 					UpperCamelCase__(    self					,						__lowerCamelCase					,						__lowerCamelCase					,						__lowerCamelCase					,						__lowerCamelCase					,						__lowerCamelCase = None					,						__lowerCamelCase = True					,						):



     '''simple docstring'''




     __A    :						Optional[int]						 =     timestep

     if key is None:
         __A    :						List[Any]						 =     jax.random.PRNGKey(0  )

     if model_output.shape[1] == sample.shape[1] * 2 and self.config.variance_type in ["learned", "learned_range"]:
         __A ,       __A    :						Tuple						 =     jnp.split(__lowerCamelCase					,						sample.shape[1]					,						axis=1  )
     else:
         __A    :						List[str]						 =     None

     # 1. compute alphas, betas
     __A    :						Dict						 =     state.common.alphas_cumprod[t]
     __A    :						int						 =     jnp.where(t > 0					,						state.common.alphas_cumprod[t - 1]					,						jnp.array(1.0					,						dtype=self.dtype  )  )
     __A    :						Tuple						 =     1 - alpha_prod_t
     __A    :						Optional[int]						 =     1 - alpha_prod_t_prev

     # 2. compute predicted original sample from predicted noise also called
     # "predicted x_0" of formula (15) from https://arxiv.org/pdf/2006.11239.pdf
     if self.config.prediction_type == "epsilon":
         __A    :						Optional[Any]						 =     (sample - beta_prod_t ** 0.5 * model_output) / alpha_prod_t ** 0.5
     elif self.config.prediction_type == "sample":
         __A    :						Any						 =     model_output
     elif self.config.prediction_type == "v_prediction":
         __A    :						str						 =     (alpha_prod_t**0.5) * sample - (beta_prod_t**0.5) * model_output
     else:
         raise ValueError(
             F"""prediction_type given as {self.config.prediction_type} must be one of `epsilon`, `sample` """
             ''' for the FlaxDDPMScheduler.'''  )

     # 3. Clip "predicted x_0"
     if self.config.clip_sample:
         __A    :						str						 =     jnp.clip(__lowerCamelCase					,						-1					,						1  )

     # 4. Compute coefficients for pred_original_sample x_0 and current sample x_t
     # See formula (7) from https://arxiv.org/pdf/2006.11239.pdf
     __A    :						Optional[Any]						 =     (alpha_prod_t_prev ** 0.5 * state.common.betas[t]) / beta_prod_t
     __A    :						Union[str, Any]						 =     state.common.alphas[t] ** 0.5 * beta_prod_t_prev / beta_prod_t

     # 5. Compute predicted previous sample µ_t
     # See formula (7) from https://arxiv.org/pdf/2006.11239.pdf
     __A    :						Optional[Any]						 =     pred_original_sample_coeff * pred_original_sample + current_sample_coeff * sample

     # 6. Add noise
     def random_variance():
         __A    :						List[Any]						 =     jax.random.split(__lowerCamelCase					,						num=1  )
         __A    :						List[str]						 =     jax.random.normal(__lowerCamelCase					,						shape=model_output.shape					,						dtype=self.dtype  )
         return (self._get_variance(__lowerCamelCase					,						__lowerCamelCase					,						predicted_variance=__lowerCamelCase  ) ** 0.5) * noise

     __A    :						Optional[Any]						 =     jnp.where(t > 0					,						random_variance()					,						jnp.zeros(model_output.shape					,						dtype=self.dtype  )  )

     __A    :						Any						 =     pred_prev_sample + variance

     if not return_dict:
         return (pred_prev_sample, state)

     return FlaxDDPMSchedulerOutput(prev_sample=__lowerCamelCase					,						state=__lowerCamelCase  )

 def 					UpperCamelCase__(    self					,						__lowerCamelCase					,						__lowerCamelCase					,						__lowerCamelCase					,						__lowerCamelCase					,						):



     '''simple docstring'''




     return add_noise_common(state.common					,						__lowerCamelCase					,						__lowerCamelCase					,						__lowerCamelCase  )

 def 					UpperCamelCase__(    self					,						__lowerCamelCase					,						__lowerCamelCase					,						__lowerCamelCase					,						__lowerCamelCase					,						):



     '''simple docstring'''




     return get_velocity_common(state.common					,						__lowerCamelCase					,						__lowerCamelCase					,						__lowerCamelCase  )







 def __len__(    self  ):



     '''simple docstring'''




     return self.config.num_train_timesteps




