





'''simple docstring'''

from typing import Optional, Tuple, Union

import flax
import flax.linen as nn
import jax
import jax.numpy as jnp
from flax.core.frozen_dict import FrozenDict

from ..configuration_utils import ConfigMixin, flax_register_to_config
from ..utils import BaseOutput
from .embeddings_flax import FlaxTimestepEmbedding, FlaxTimesteps
from .modeling_flax_utils import FlaxModelMixin
from .unet_ad_blocks_flax import (
    FlaxCrossAttnDownBlockaD,
    FlaxDownBlockaD,
    FlaxUNetMidBlockaDCrossAttn,
)






@flax.struct.dataclass
class        UpperCAmelCase     (	UpperCamelCase__    ):
 __lowercase       		=    42
 __lowercase       		=    42






class        UpperCAmelCase     (	nn.Module    ):
 __lowercase       		=    42
 __lowercase       		=    (16, 32, 96, 256)
 __lowercase       		=    jnp.floataa
 def 					UpperCAmelCase_							(    self				:Optional[Any]							)->							List[Any]:
     A__    =      nn.Conv(
         self.block_out_channels[0]						,	kernel_size=(3, 3)						,	padding=((1, 1), (1, 1))						,	dtype=self.dtype						,	)

     A__    =      []
     for i in range(len(self.block_out_channels							) - 1							):
         A__    =      self.block_out_channels[i]
         A__    =      self.block_out_channels[i + 1]
         A__    =      nn.Conv(
             lowercase_						,	kernel_size=(3, 3)						,	padding=((1, 1), (1, 1))						,	dtype=self.dtype						,	)
         blocks.append(lowercase_							)
         A__    =      nn.Conv(
             lowercase_						,	kernel_size=(3, 3)						,	strides=(2, 2)						,	padding=((1, 1), (1, 1))						,	dtype=self.dtype						,	)
         blocks.append(lowercase_							)
     A__    =      blocks

     A__    =      nn.Conv(
         self.conditioning_embedding_channels						,	kernel_size=(3, 3)						,	padding=((1, 1), (1, 1))						,	kernel_init=nn.initializers.zeros_init()						,	bias_init=nn.initializers.zeros_init()						,	dtype=self.dtype						,	)



 def __call__(    self				:Optional[Any]						,	lowercase_				:Any							)->							Dict:
     A__    =      self.conv_in(lowercase_							)
     A__    =      nn.silu(lowercase_							)

     for block in self.blocks:
         A__    =      block(lowercase_							)
         A__    =      nn.silu(lowercase_							)

     A__    =      self.conv_out(lowercase_							)

     return embedding






@flax_register_to_config
class        UpperCAmelCase     (	nn.Module      ,	UpperCamelCase__      ,	UpperCamelCase__    ):
 __lowercase       		=    32
 __lowercase       		=    4
 __lowercase       		=    (
     "CrossAttnDownBlock2D",
     "CrossAttnDownBlock2D",
     "CrossAttnDownBlock2D",
     "DownBlock2D",
 )
 __lowercase       		=    False
 __lowercase       		=    (320, 640, 1280, 1280)
 __lowercase       		=    2
 __lowercase       		=    8
 __lowercase       		=    None
 __lowercase       		=    1280
 __lowercase       		=    0.0
 __lowercase       		=    False
 __lowercase       		=    jnp.floataa
 __lowercase       		=    True
 __lowercase       		=    0
 __lowercase       		=    "rgb"
 __lowercase       		=    (16, 32, 96, 256)
 def 					UpperCAmelCase_							(    self				:int						,	lowercase_				:jax.random.KeyArray							)->							FrozenDict:
     # init input tensors
     A__    =      (1, self.in_channels, self.sample_size, self.sample_size)
     A__    =      jnp.zeros(lowercase_						,	dtype=jnp.floataa							)
     A__    =      jnp.ones((1,)						,	dtype=jnp.intaa							)
     A__    =      jnp.zeros((1, 1, self.cross_attention_dim)						,	dtype=jnp.floataa							)
     A__    =      (1, 3, self.sample_size * 8, self.sample_size * 8)
     A__    =      jnp.zeros(lowercase_						,	dtype=jnp.floataa							)

     A__,     A__    =      jax.random.split(lowercase_							)
     A__    =      {"params": params_rng, "dropout": dropout_rng}

     return self.init(lowercase_						,	lowercase_						,	lowercase_						,	lowercase_						,	lowercase_							)["params"]
 def 					UpperCAmelCase_							(    self				:Optional[Any]							)->							int:
     A__    =      self.block_out_channels
     A__    =      block_out_channels[0] * 4

     # If `num_attention_heads` is not defined (which is the case for most models)
     # it will default to `attention_head_dim`. This looks weird upon first reading it and it is.
     # The reason for this behavior is to correct for incorrectly named variables that were introduced
     # when this library was created. The incorrect naming was only discovered much later in https://github.com/huggingface/diffusers/issues/2011#issuecomment-1547958131
     # Changing `attention_head_dim` to `num_attention_heads` for 40,000+ configurations is too backwards breaking
     # which is why we correct for the naming here.
     A__    =      self.num_attention_heads or self.attention_head_dim

     # input
     A__    =      nn.Conv(
         block_out_channels[0]						,	kernel_size=(3, 3)						,	strides=(1, 1)						,	padding=((1, 1), (1, 1))						,	dtype=self.dtype						,	)

     # time
     A__    =      FlaxTimesteps(
         block_out_channels[0]						,	flip_sin_to_cos=self.flip_sin_to_cos						,	freq_shift=self.config.freq_shift							)
     A__    =      FlaxTimestepEmbedding(lowercase_						,	dtype=self.dtype							)

     A__    =      FlaxControlNetConditioningEmbedding(
         conditioning_embedding_channels=block_out_channels[0]						,	block_out_channels=self.conditioning_embedding_out_channels						,	)

     A__    =      self.only_cross_attention
     if isinstance(lowercase_						,	lowercase_							):
         A__    =      (only_cross_attention,) * len(self.down_block_types							)

     if isinstance(lowercase_						,	lowercase_							):
         A__    =      (num_attention_heads,) * len(self.down_block_types							)

     # down
     A__    =      []
     A__    =      []

     A__    =      block_out_channels[0]

     A__    =      nn.Conv(
         lowercase_						,	kernel_size=(1, 1)						,	padding="VALID"						,	kernel_init=nn.initializers.zeros_init()						,	bias_init=nn.initializers.zeros_init()						,	dtype=self.dtype						,	)
     controlnet_down_blocks.append(lowercase_							)

     for i, down_block_type in enumerate(self.down_block_types							):
         A__    =      output_channel
         A__    =      block_out_channels[i]
         A__    =      i == len(lowercase_							) - 1

         if down_block_type == "CrossAttnDownBlock2D":
             A__    =      FlaxCrossAttnDownBlockaD(
                 in_channels=lowercase_						,	out_channels=lowercase_						,	dropout=self.dropout						,	num_layers=self.layers_per_block						,	num_attention_heads=num_attention_heads[i]						,	add_downsample=not is_final_block						,	use_linear_projection=self.use_linear_projection						,	only_cross_attention=only_cross_attention[i]						,	dtype=self.dtype						,	)
         else:
             A__    =      FlaxDownBlockaD(
                 in_channels=lowercase_						,	out_channels=lowercase_						,	dropout=self.dropout						,	num_layers=self.layers_per_block						,	add_downsample=not is_final_block						,	dtype=self.dtype						,	)

         down_blocks.append(lowercase_							)

         for _ in range(self.layers_per_block							):
             A__    =      nn.Conv(
                 lowercase_						,	kernel_size=(1, 1)						,	padding="VALID"						,	kernel_init=nn.initializers.zeros_init()						,	bias_init=nn.initializers.zeros_init()						,	dtype=self.dtype						,	)
             controlnet_down_blocks.append(lowercase_							)

         if not is_final_block:
             A__    =      nn.Conv(
                 lowercase_						,	kernel_size=(1, 1)						,	padding="VALID"						,	kernel_init=nn.initializers.zeros_init()						,	bias_init=nn.initializers.zeros_init()						,	dtype=self.dtype						,	)
             controlnet_down_blocks.append(lowercase_							)

     A__    =      down_blocks
     A__    =      controlnet_down_blocks

     # mid
     A__    =      block_out_channels[-1]
     A__    =      FlaxUNetMidBlockaDCrossAttn(
         in_channels=lowercase_						,	dropout=self.dropout						,	num_attention_heads=num_attention_heads[-1]						,	use_linear_projection=self.use_linear_projection						,	dtype=self.dtype						,	)

     A__    =      nn.Conv(
         lowercase_						,	kernel_size=(1, 1)						,	padding="VALID"						,	kernel_init=nn.initializers.zeros_init()						,	bias_init=nn.initializers.zeros_init()						,	dtype=self.dtype						,	)



 def __call__(    self				:Tuple						,	lowercase_				:Optional[int]						,	lowercase_				:List[Any]						,	lowercase_				:List[Any]						,	lowercase_				:Union[str, Any]						,	lowercase_				:float = 1.0						,	lowercase_				:bool = True						,	lowercase_				:bool = False						,	)->							Union[FlaxControlNetOutput, Tuple]:
     A__    =      self.controlnet_conditioning_channel_order
     if channel_order == "bgr":
         A__    =      jnp.flip(lowercase_						,	axis=1							)

     # 1. time
     if not isinstance(lowercase_						,	jnp.ndarray							):
         A__    =      jnp.array([timesteps]						,	dtype=jnp.intaa							)
     elif isinstance(lowercase_						,	jnp.ndarray							) and len(timesteps.shape							) == 0:
         A__    =      timesteps.astype(dtype=jnp.floataa							)
         A__    =      jnp.expand_dims(lowercase_						,	0							)

     A__    =      self.time_proj(lowercase_							)
     A__    =      self.time_embedding(lowercase_							)

     # 2. pre-process
     A__    =      jnp.transpose(lowercase_						,	(0, 2, 3, 1)							)
     A__    =      self.conv_in(lowercase_							)

     A__    =      jnp.transpose(lowercase_						,	(0, 2, 3, 1)							)
     A__    =      self.controlnet_cond_embedding(lowercase_							)
     sample += controlnet_cond

     # 3. down
     A__    =      (sample,)
     for down_block in self.down_blocks:
         if isinstance(lowercase_						,	lowercase_							):
             A__,     A__    =      down_block(lowercase_						,	lowercase_						,	lowercase_						,	deterministic=not train							)
         else:
             A__,     A__    =      down_block(lowercase_						,	lowercase_						,	deterministic=not train							)
         down_block_res_samples += res_samples

     # 4. mid
     A__    =      self.mid_block(lowercase_						,	lowercase_						,	lowercase_						,	deterministic=not train							)

     # 5. contronet blocks
     A__    =      ()
     for down_block_res_sample, controlnet_block in zip(lowercase_						,	self.controlnet_down_blocks							):
         A__    =      controlnet_block(lowercase_							)
         controlnet_down_block_res_samples += (down_block_res_sample,)

     A__    =      controlnet_down_block_res_samples

     A__    =      self.controlnet_mid_block(lowercase_							)

     # 6. scaling
     A__    =      [sample * conditioning_scale for sample in down_block_res_samples]
     mid_block_res_sample *= conditioning_scale

     if not return_dict:
         return (down_block_res_samples, mid_block_res_sample)

     return FlaxControlNetOutput(
         down_block_res_samples=lowercase_						,	mid_block_res_sample=lowercase_							)










'''simple docstring'''

from __future__ import annotations

import unittest

from transformers import DistilBertConfig, is_tf_available
from transformers.testing_utils import require_tf, slow

from ...test_configuration_common import ConfigTester
from ...test_modeling_tf_common import TFModelTesterMixin, ids_tensor, random_attention_mask
from ...test_pipeline_mixin import PipelineTesterMixin


if is_tf_available():
  import tensorflow as tf

  from transformers.models.distilbert.modeling_tf_distilbert import (
      TF_DISTILBERT_PRETRAINED_MODEL_ARCHIVE_LIST,
      TFDistilBertForMaskedLM,
      TFDistilBertForMultipleChoice,
      TFDistilBertForQuestionAnswering,
      TFDistilBertForSequenceClassification,
      TFDistilBertForTokenClassification,
      TFDistilBertModel,
  )






class        UpperCAmelCase     :
 def __init__(    self				:str						,	lowercase_				:str						,	)->							str:
     A__    =      parent
     A__    =      13
     A__    =      7
     A__    =      True
     A__    =      True
     A__    =      False
     A__    =      True
     A__    =      99
     A__    =      32
     A__    =      2
     A__    =      4
     A__    =      37
     A__    =      "gelu"
     A__    =      0.1
     A__    =      0.1
     A__    =      5_12
     A__    =      16
     A__    =      2
     A__    =      0.0_2
     A__    =      3
     A__    =      4
     A__    =      None
 def 					UpperCAmelCase_							(    self				:Union[str, Any]							)->							int:
     A__    =      ids_tensor([self.batch_size, self.seq_length]						,	self.vocab_size							)

     A__    =      None
     if self.use_input_mask:
         A__    =      random_attention_mask([self.batch_size, self.seq_length]							)

     A__    =      None
     A__    =      None
     A__    =      None
     if self.use_labels:
         A__    =      ids_tensor([self.batch_size]						,	self.type_sequence_label_size							)
         A__    =      ids_tensor([self.batch_size, self.seq_length]						,	self.num_labels							)
         A__    =      ids_tensor([self.batch_size]						,	self.num_choices							)

     A__    =      DistilBertConfig(
         vocab_size=self.vocab_size						,	dim=self.hidden_size						,	n_layers=self.num_hidden_layers						,	n_heads=self.num_attention_heads						,	hidden_dim=self.intermediate_size						,	hidden_act=self.hidden_act						,	dropout=self.hidden_dropout_prob						,	attention_dropout=self.attention_probs_dropout_prob						,	max_position_embeddings=self.max_position_embeddings						,	initializer_range=self.initializer_range						,	)

     return config, input_ids, input_mask, sequence_labels, token_labels, choice_labels
 def 					UpperCAmelCase_							(    self				:str						,	lowercase_				:Optional[int]						,	lowercase_				:List[str]						,	lowercase_				:Any						,	lowercase_				:Union[str, Any]						,	lowercase_				:Optional[int]						,	lowercase_				:str							)->							List[str]:
     A__    =      TFDistilBertModel(config=lowercase_							)
     A__    =      {"input_ids": input_ids, "attention_mask": input_mask}

     A__    =      model(lowercase_							)

     A__    =      [input_ids, input_mask]

     A__    =      model(lowercase_							)

     self.parent.assertEqual(result.last_hidden_state.shape						,	(self.batch_size, self.seq_length, self.hidden_size)							)
 def 					UpperCAmelCase_							(    self				:List[str]						,	lowercase_				:str						,	lowercase_				:Optional[Any]						,	lowercase_				:Optional[int]						,	lowercase_				:Optional[int]						,	lowercase_				:Optional[int]						,	lowercase_				:Union[str, Any]							)->							Optional[int]:
     A__    =      TFDistilBertForMaskedLM(config=lowercase_							)
     A__    =      {"input_ids": input_ids, "attention_mask": input_mask}
     A__    =      model(lowercase_							)
     self.parent.assertEqual(result.logits.shape						,	(self.batch_size, self.seq_length, self.vocab_size)							)
 def 					UpperCAmelCase_							(    self				:Any						,	lowercase_				:str						,	lowercase_				:str						,	lowercase_				:Optional[int]						,	lowercase_				:str						,	lowercase_				:List[Any]						,	lowercase_				:Union[str, Any]							)->							Optional[int]:
     A__    =      TFDistilBertForQuestionAnswering(config=lowercase_							)
     A__    =      {
         "input_ids": input_ids,
         "attention_mask": input_mask,
     }
     A__    =      model(lowercase_							)
     self.parent.assertEqual(result.start_logits.shape						,	(self.batch_size, self.seq_length)							)
     self.parent.assertEqual(result.end_logits.shape						,	(self.batch_size, self.seq_length)							)
 def 					UpperCAmelCase_							(    self				:Union[str, Any]						,	lowercase_				:Optional[int]						,	lowercase_				:Any						,	lowercase_				:Dict						,	lowercase_				:Tuple						,	lowercase_				:Optional[Any]						,	lowercase_				:Optional[int]							)->							Any:
     A__    =      self.num_labels
     A__    =      TFDistilBertForSequenceClassification(lowercase_							)
     A__    =      {"input_ids": input_ids, "attention_mask": input_mask}
     A__    =      model(lowercase_							)
     self.parent.assertEqual(result.logits.shape						,	(self.batch_size, self.num_labels)							)
 def 					UpperCAmelCase_							(    self				:str						,	lowercase_				:Optional[Any]						,	lowercase_				:List[Any]						,	lowercase_				:Dict						,	lowercase_				:Tuple						,	lowercase_				:int						,	lowercase_				:Union[str, Any]							)->							str:
     A__    =      self.num_choices
     A__    =      TFDistilBertForMultipleChoice(lowercase_							)
     A__    =      tf.tile(tf.expand_dims(lowercase_						,	1							)						,	(1, self.num_choices, 1)							)
     A__    =      tf.tile(tf.expand_dims(lowercase_						,	1							)						,	(1, self.num_choices, 1)							)
     A__    =      {
         "input_ids": multiple_choice_inputs_ids,
         "attention_mask": multiple_choice_input_mask,
     }
     A__    =      model(lowercase_							)
     self.parent.assertEqual(result.logits.shape						,	(self.batch_size, self.num_choices)							)
 def 					UpperCAmelCase_							(    self				:str						,	lowercase_				:Any						,	lowercase_				:List[str]						,	lowercase_				:Any						,	lowercase_				:int						,	lowercase_				:List[Any]						,	lowercase_				:Tuple							)->							Tuple:
     A__    =      self.num_labels
     A__    =      TFDistilBertForTokenClassification(lowercase_							)
     A__    =      {"input_ids": input_ids, "attention_mask": input_mask}
     A__    =      model(lowercase_							)
     self.parent.assertEqual(result.logits.shape						,	(self.batch_size, self.seq_length, self.num_labels)							)



 def 					UpperCAmelCase_							(    self				:Any							)->							Union[str, Any]:
     A__    =      self.prepare_config_and_inputs()
     ((A__),     (A__),     (A__),     (A__),     (A__),     (A__))    =      config_and_inputs
     A__    =      {"input_ids": input_ids, "attention_mask": input_mask}
     return config, inputs_dict






@require_tf
class        UpperCAmelCase     (	UpperCamelCase__      ,	UpperCamelCase__      ,	unittest.TestCase    ):
 __lowercase       		=    (
     (
         TFDistilBertModel,
         TFDistilBertForMaskedLM,
         TFDistilBertForQuestionAnswering,
         TFDistilBertForSequenceClassification,
         TFDistilBertForTokenClassification,
         TFDistilBertForMultipleChoice,
     )
     if is_tf_available()
     else None
 )
 __lowercase       		=    (
     {
         """feature-extraction""": TFDistilBertModel,
         """fill-mask""": TFDistilBertForMaskedLM,
         """question-answering""": TFDistilBertForQuestionAnswering,
         """text-classification""": TFDistilBertForSequenceClassification,
         """token-classification""": TFDistilBertForTokenClassification,
         """zero-shot""": TFDistilBertForSequenceClassification,
     }
     if is_tf_available()
     else {}
 )
 __lowercase       		=    False
 __lowercase       		=    False
 def 					UpperCAmelCase_							(    self				:Optional[Any]							)->							List[Any]:
     A__    =      TFDistilBertModelTester(self							)
     A__    =      ConfigTester(self						,	config_class=lowercase_						,	dim=37							)
 def 					UpperCAmelCase_							(    self				:Tuple							)->							Tuple:
     self.config_tester.run_common_tests()
 def 					UpperCAmelCase_							(    self				:int							)->							Tuple:
     A__    =      self.model_tester.prepare_config_and_inputs()
     self.model_tester.create_and_check_distilbert_model(*lowercase_							)
 def 					UpperCAmelCase_							(    self				:Optional[int]							)->							Optional[Any]:
     A__    =      self.model_tester.prepare_config_and_inputs()
     self.model_tester.create_and_check_distilbert_for_masked_lm(*lowercase_							)
 def 					UpperCAmelCase_							(    self				:str							)->							str:
     A__    =      self.model_tester.prepare_config_and_inputs()
     self.model_tester.create_and_check_distilbert_for_question_answering(*lowercase_							)
 def 					UpperCAmelCase_							(    self				:List[str]							)->							Dict:
     A__    =      self.model_tester.prepare_config_and_inputs()
     self.model_tester.create_and_check_distilbert_for_sequence_classification(*lowercase_							)
 def 					UpperCAmelCase_							(    self				:List[str]							)->							Optional[int]:
     A__    =      self.model_tester.prepare_config_and_inputs()
     self.model_tester.create_and_check_distilbert_for_multiple_choice(*lowercase_							)
 def 					UpperCAmelCase_							(    self				:str							)->							int:
     A__    =      self.model_tester.prepare_config_and_inputs()
     self.model_tester.create_and_check_distilbert_for_token_classification(*lowercase_							)



 @slow
 def 					UpperCAmelCase_							(    self				:List[str]							)->							Dict:
     for model_name in list(TF_DISTILBERT_PRETRAINED_MODEL_ARCHIVE_LIST[:1]							):
         A__    =      TFDistilBertModel.from_pretrained(lowercase_							)
         self.assertIsNotNone(lowercase_							)






@require_tf
class        UpperCAmelCase     (	unittest.TestCase    ):
 @slow
 def 					UpperCAmelCase_							(    self				:List[Any]							)->							Any:
     A__    =      TFDistilBertModel.from_pretrained("distilbert-base-uncased"							)
     A__    =      tf.constant([[0, 1, 2, 3, 4, 5]]							)
     A__    =      model(lowercase_							)[0]

     A__    =      [1, 6, 7_68]
     self.assertEqual(output.shape						,	lowercase_							)

     A__    =      tf.constant(
         [
             [
                 [0.1_9_2_6_1_8_8_5, -0.1_3_7_3_2_9_5_5, 0.4_1_1_9_7_9_9],
                 [0.2_2_1_5_0_1_5_6, -0.0_7_4_2_2_6_6_1, 0.3_9_0_3_7_2_0_4],
                 [0.2_2_7_5_6_0_1_8, -0.0_8_9_6_4_1_4, 0.3_7_0_1_4_6_7],
             ]
         ]							)
     tf.debugging.assert_near(output[:, :3, :3]						,	lowercase_						,	atol=1E-4							)



