



"""simple docstring"""






from typing import Optional, Tuple

import jax
import jax.numpy as jnp
from flax import linen as nn
from flax.core.frozen_dict import FrozenDict
from transformers import CLIPConfig, FlaxPreTrainedModel
from transformers.models.clip.modeling_flax_clip import FlaxCLIPVisionModule


def     _lowercase       (						__snake_case      ,__snake_case      ,__snake_case=1e-1_2			):
			__lowerCAmelCase   :						List[str]				= jnp.divide(emb_a.T      ,jnp.clip(jnp.linalg.norm(__snake_case      ,axis=1			)      ,a_min=__snake_case			)			).T
			__lowerCAmelCase   :						Tuple				= jnp.divide(emb_a.T      ,jnp.clip(jnp.linalg.norm(__snake_case      ,axis=1			)      ,a_min=__snake_case			)			).T
			return jnp.matmul(__snake_case      ,norm_emb_a.T			)



class     A__     (						nn.Module							):




	'''simple docstring'''





	SCREAMING_SNAKE_CASE					   =			4_2
	SCREAMING_SNAKE_CASE					   =			jnp.floataa

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

				"""simple docstring"""

				__lowerCAmelCase   :						Any				= FlaxCLIPVisionModule(self.config.vision_config)
				__lowerCAmelCase   :						Optional[Any]				= nn.Dense(self.config.projection_dim ,    use_bias=_SCREAMING_SNAKE_CASE ,    dtype=self.dtype)

				__lowerCAmelCase   :						int				= self.param("concept_embeds" ,    jax.nn.initializers.ones ,    (17, self.config.projection_dim))
				__lowerCAmelCase   :						List[Any]				= self.param(
				    "special_care_embeds" ,    jax.nn.initializers.ones ,    (3, self.config.projection_dim))

				__lowerCAmelCase   :						Optional[int]				= self.param("concept_embeds_weights" ,    jax.nn.initializers.ones ,    (17,))
				__lowerCAmelCase   :						Dict				= self.param("special_care_embeds_weights" ,    jax.nn.initializers.ones ,    (3,))







	def __call__(  self:	List[Any] ,    _SCREAMING_SNAKE_CASE:	Any)       -> int:

				"""simple docstring"""

				__lowerCAmelCase   :						Tuple				= self.vision_model(_SCREAMING_SNAKE_CASE)[1]
				__lowerCAmelCase   :						Tuple				= self.visual_projection(_SCREAMING_SNAKE_CASE)

				__lowerCAmelCase   :						int				= jax_cosine_distance(_SCREAMING_SNAKE_CASE ,    self.special_care_embeds)
				__lowerCAmelCase   :						Any				= jax_cosine_distance(_SCREAMING_SNAKE_CASE ,    self.concept_embeds)

				# increase this value to create a stronger `nfsw` filter
				# at the cost of increasing the possibility of filtering benign image inputs
				__lowerCAmelCase   :						Tuple				= 0.0

				__lowerCAmelCase   :						int				= special_cos_dist - self.special_care_embeds_weights[None, :] + adjustment
				__lowerCAmelCase   :						Optional[Any]				= jnp.round(_SCREAMING_SNAKE_CASE ,    3)
				__lowerCAmelCase   :						Optional[int]				= jnp.any(special_scores > 0 ,    axis=1 ,    keepdims=_SCREAMING_SNAKE_CASE)
				# Use a lower threshold if an image has any special care concept
				__lowerCAmelCase   :						List[str]				= is_special_care * 0.01

				__lowerCAmelCase   :						Optional[Any]				= cos_dist - self.concept_embeds_weights[None, :] + special_adjustment
				__lowerCAmelCase   :						Optional[Any]				= jnp.round(_SCREAMING_SNAKE_CASE ,    3)
				__lowerCAmelCase   :						int				= jnp.any(concept_scores > 0 ,    axis=1)

				return has_nsfw_concepts






class     A__     (						__SCREAMING_SNAKE_CASE							):




	'''simple docstring'''





	SCREAMING_SNAKE_CASE					   =			CLIPConfig
	SCREAMING_SNAKE_CASE					   =			'clip_input'
	SCREAMING_SNAKE_CASE					   =			FlaxStableDiffusionSafetyCheckerModule

	def __init__(  self:	int ,    _SCREAMING_SNAKE_CASE:	CLIPConfig ,    _SCREAMING_SNAKE_CASE:	Optional[Tuple] = None ,    _SCREAMING_SNAKE_CASE:	int = 0 ,    _SCREAMING_SNAKE_CASE:	jnp.dtype = jnp.floataa ,    _SCREAMING_SNAKE_CASE:	bool = True ,    **_SCREAMING_SNAKE_CASE:	List[Any] ,    )       -> Union[str, Any]:

				"""simple docstring"""

				if input_shape is None:
							__lowerCAmelCase   :						List[str]				= (1, 224, 224, 3)
				__lowerCAmelCase   :						str				= self.module_class(config=_SCREAMING_SNAKE_CASE ,    dtype=_SCREAMING_SNAKE_CASE ,    **_SCREAMING_SNAKE_CASE)
				super().__init__(_SCREAMING_SNAKE_CASE ,    _SCREAMING_SNAKE_CASE ,    input_shape=_SCREAMING_SNAKE_CASE ,    seed=_SCREAMING_SNAKE_CASE ,    dtype=_SCREAMING_SNAKE_CASE ,    _do_init=_do_init)

	def 							_SCREAMING_SNAKE_CASE     (  self:	Dict ,    _SCREAMING_SNAKE_CASE:	jax.random.KeyArray ,    _SCREAMING_SNAKE_CASE:	Tuple ,    _SCREAMING_SNAKE_CASE:	FrozenDict = None)       -> FrozenDict:

				"""simple docstring"""

				__lowerCAmelCase   :						List[Any]				= jax.random.normal(_SCREAMING_SNAKE_CASE ,    _SCREAMING_SNAKE_CASE)

				__lowerCAmelCase   :						Any				= jax.random.split(_SCREAMING_SNAKE_CASE)
				__lowerCAmelCase   :						Optional[int]				= {"params": params_rng, "dropout": dropout_rng}

				__lowerCAmelCase   :						Dict				= self.module.init(_SCREAMING_SNAKE_CASE ,    _SCREAMING_SNAKE_CASE)["params"]

				return random_params







	def __call__(  self:	Tuple ,    _SCREAMING_SNAKE_CASE:	List[Any] ,    _SCREAMING_SNAKE_CASE:	dict = None ,    )       -> Tuple:

				"""simple docstring"""

				__lowerCAmelCase   :						Dict				= jnp.transpose(_SCREAMING_SNAKE_CASE ,    (0, 2, 3, 1))

				return self.module.apply(
				    {"params": params or self.params} ,    jnp.array(_SCREAMING_SNAKE_CASE ,    dtype=jnp.floataa) ,    rngs={} ,    )




"""simple docstring"""






import math


def     _lowercase       (						__snake_case      ,__snake_case			)		->     float:
			return math.pow(__snake_case      ,2			) - a


def     _lowercase       (						__snake_case			)		->     float:
			return 2 * x


def     _lowercase       (						__snake_case			)		->     float:
			__lowerCAmelCase   :						List[str]				= 2.0

			while start <= a:
						__lowerCAmelCase   :						Union[str, Any]				= math.pow(__snake_case      ,2			)

			return start


def     _lowercase       (						__snake_case      ,__snake_case = 9_999      ,__snake_case = 0.00000000000001			)		->     float:

			if a < 0:
						raise ValueError("math domain error"			)

			__lowerCAmelCase   :						List[Any]				= get_initial_point(__snake_case			)

			for _ in range(__snake_case			):
						__lowerCAmelCase   :						Union[str, Any]				= value
						__lowerCAmelCase   :						List[str]				= value - fx(__snake_case      ,__snake_case			) / fx_derivative(__snake_case			)
						if abs(prev_value - value			) < tolerance:
									return value

			return value


if __name__ == "__main__":
	from doctest import testmod

	testmod()