






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 					lowerCamelCase__					(_UpperCAmelCase				,							_UpperCAmelCase				,							_UpperCAmelCase=1e-12):
			SCREAMING_SNAKE_CASE							=			jnp.divide(emb_a.T				,							jnp.clip(jnp.linalg.norm(_UpperCAmelCase				,							axis=1)				,							a_min=_UpperCAmelCase)).T
			SCREAMING_SNAKE_CASE							=			jnp.divide(emb_a.T				,							jnp.clip(jnp.linalg.norm(_UpperCAmelCase				,							axis=1)				,							a_min=_UpperCAmelCase)).T
			return jnp.matmul(_UpperCAmelCase				,							norm_emb_a.T)

class 		_snake_case	(	nn.Module       ):
			_lowercase		:     CLIPConfig
			_lowercase		:     jnp.dtype        =				jnp.floataa






			def  SCREAMING_SNAKE_CASE__     (							self)   ->							Optional[int]:
						SCREAMING_SNAKE_CASE							=			FlaxCLIPVisionModule(self.config.vision_config)
						SCREAMING_SNAKE_CASE							=			nn.Dense(self.config.projection_dim		,       use_bias=a		,       dtype=self.dtype)

						SCREAMING_SNAKE_CASE							=			self.param('concept_embeds'		,       jax.nn.initializers.ones		,       (17, self.config.projection_dim))
						SCREAMING_SNAKE_CASE							=			self.param(
						    'special_care_embeds'		,       jax.nn.initializers.ones		,       (3, self.config.projection_dim))

						SCREAMING_SNAKE_CASE							=			self.param('concept_embeds_weights'		,       jax.nn.initializers.ones		,       (17,))
						SCREAMING_SNAKE_CASE							=			self.param('special_care_embeds_weights'		,       jax.nn.initializers.ones		,       (3,))







			def __call__(							self		,       a)   ->							Tuple:
						SCREAMING_SNAKE_CASE							=			self.vision_model(a)[1]
						SCREAMING_SNAKE_CASE							=			self.visual_projection(a)

						SCREAMING_SNAKE_CASE							=			jax_cosine_distance(a		,       self.special_care_embeds)
						SCREAMING_SNAKE_CASE							=			jax_cosine_distance(a		,       self.concept_embeds)

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

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

						SCREAMING_SNAKE_CASE							=			cos_dist - self.concept_embeds_weights[None, :] + special_adjustment
						SCREAMING_SNAKE_CASE							=			jnp.round(a		,       3)
						SCREAMING_SNAKE_CASE							=			jnp.any(concept_scores > 0		,       axis=1)

						return has_nsfw_concepts



class 		_snake_case	(	A__       ):
			_lowercase		:     Optional[Any]        =				CLIPConfig
			_lowercase		:     Optional[Any]        =				'''clip_input'''
			_lowercase		:     Dict        =				FlaxStableDiffusionSafetyCheckerModule






			def __init__(							self		,       a		,       a = None		,       a = 0		,       a = jnp.floataa		,       a = True		,       **a		,       )   ->							Dict:
						if input_shape is None:
									SCREAMING_SNAKE_CASE							=			(1, 224, 224, 3)
						SCREAMING_SNAKE_CASE							=			self.module_class(config=a		,       dtype=a		,       **a)
						super().__init__(a		,       a		,       input_shape=a		,       seed=a		,       dtype=a		,       _do_init=_do_init)






			def  SCREAMING_SNAKE_CASE__     (							self		,       a		,       a		,       a = None)   ->							FrozenDict:
						# init input tensor
						SCREAMING_SNAKE_CASE							=			jax.random.normal(a		,       a)

						SCREAMING_SNAKE_CASE					,    SCREAMING_SNAKE_CASE							=			jax.random.split(a)
						SCREAMING_SNAKE_CASE							=			{'params': params_rng, 'dropout': dropout_rng}

						SCREAMING_SNAKE_CASE							=			self.module.init(a		,       a)['params']

						return random_params







			def __call__(							self		,       a		,       a = None		,       )   ->							Optional[int]:
						SCREAMING_SNAKE_CASE							=			jnp.transpose(a		,       (0, 2, 3, 1))

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






from decimal import Decimal, getcontext
from math import ceil, factorial

def  __lowerCamelCase					( _lowercase			)       ->      str:

     if not isinstance(_lowercase						,     _lowercase			):
          raise TypeError('Undefined for non-integers'			)
     elif precision < 1:
          raise ValueError('Undefined for non-natural numbers'			)

     UpperCamelCase						    =					precision
     UpperCamelCase						    =					ceil(precision / 14			)
     UpperCamelCase						    =					426880 * Decimal(10005			).sqrt()
     UpperCamelCase						    =					1
     UpperCamelCase						    =					13591409
     UpperCamelCase						    =					Decimal(_lowercase			)
     for k in range(1						,     _lowercase			):
          UpperCamelCase						    =					factorial(6 * k			) // (factorial(3 * k			) * factorial(_lowercase			) ** 3)
          linear_term += 545140134
          exponential_term *= -262537412640768000
          partial_sum += Decimal(multinomial_term * linear_term			) / exponential_term
     return str(constant_term / partial_sum			)[:-1]


if __name__ == "__main__":
 _snake_case      = 50
 print(F"The first {n} digits of pi is: {pi(n)}")


