



'''simple docstring'''





from __future__ import annotations






class    lowerCamelCase_       :
  """simple docstring"""




  def __init__(					self :      str     ,      _a :      int  )   ->   None:
    __lowerCamelCase			:			Tuple       			=					order

    # a_{0} ... a_{k}
    __lowerCamelCase			:			str       			=					[1.0] + [0.0] * order
    # b_{0} ... b_{k}
    __lowerCamelCase			:			Tuple       			=					[1.0] + [0.0] * order

    # x[n-1] ... x[n-k]
    __lowerCamelCase			:			Any       			=					[0.0] * self.order
    # y[n-1] ... y[n-k]
    __lowerCamelCase			:			List[Any]       			=					[0.0] * self.order




  def 		_lowercase	(					self :      int     ,      _a :      list[float]     ,      _a :      list[float]  )   ->   None:
    if len(_a  ) < self.order:
      __lowerCamelCase			:			List[str]       			=					[1.0, *a_coeffs]

    if len(_a  ) != self.order + 1:
      __lowerCamelCase			:			List[str]       			=					(
          f'Expected a_coeffs to have {self.order + 1} elements '
          f'for {self.order}-order filter, got {len(_a  )}'
      )
      raise ValueError(_a  )

    if len(_a  ) != self.order + 1:
      __lowerCamelCase			:			List[str]       			=					(
          f'Expected b_coeffs to have {self.order + 1} elements '
          f'for {self.order}-order filter, got {len(_a  )}'
      )
      raise ValueError(_a  )

    __lowerCamelCase			:			Union[str, Any]       			=					a_coeffs
    __lowerCamelCase			:			List[str]       			=					b_coeffs




  def 		_lowercase	(					self :      Tuple     ,      _a :      float  )   ->   float:
    __lowerCamelCase			:			Any       			=					0.0

    # Start at index 1 and do index 0 at the end.
    for i in range(1     ,      self.order + 1  ):
      result += (
          self.b_coeffs[i] * self.input_history[i - 1]
          - self.a_coeffs[i] * self.output_history[i - 1]
      )

    __lowerCamelCase			:			Union[str, Any]       			=					(result + self.b_coeffs[0] * sample) / self.a_coeffs[0]

    __lowerCamelCase			:			Optional[Any]       			=					self.input_history[:-1]
    __lowerCamelCase			:			Tuple       			=					self.output_history[:-1]

    __lowerCamelCase			:			Dict       			=					sample
    __lowerCamelCase			:			Optional[Any]       			=					result

    return result











'''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 			a_	(				_lowerCAmelCase			,_lowerCAmelCase			,_lowerCAmelCase=1E-12			) -> List[Any]:
  __lowerCamelCase			:			int       			=					jnp.divide(emb_a.T			,jnp.clip(jnp.linalg.norm(_lowerCAmelCase			,axis=1			)			,a_min=_lowerCAmelCase			)			).T
  __lowerCamelCase			:			Tuple       			=					jnp.divide(emb_a.T			,jnp.clip(jnp.linalg.norm(_lowerCAmelCase			,axis=1			)			,a_min=_lowerCAmelCase			)			).T
  return jnp.matmul(_lowerCAmelCase			,norm_emb_a.T			)






class    lowerCamelCase_       (							nn.Module	):
  """simple docstring"""




  a_				   =42
  a_				   =jnp.floataa




  def 		_lowercase	(					self :      Optional[Any]  )   ->   Tuple:
    __lowerCamelCase			:			Optional[int]       			=					FlaxCLIPVisionModule(self.config.vision_config  )
    __lowerCamelCase			:			Any       			=					nn.Dense(self.config.projection_dim     ,      use_bias=_a     ,      dtype=self.dtype  )

    __lowerCamelCase			:			str       			=					self.param('concept_embeds'     ,      jax.nn.initializers.ones     ,      (17, self.config.projection_dim)  )
    __lowerCamelCase			:			str       			=					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 :      Tuple     ,      _a :      Any  )   ->   Optional[int]:
    __lowerCamelCase			:			int       			=					self.vision_model(_a  )[1]
    __lowerCamelCase			:			Optional[int]       			=					self.visual_projection(_a  )

    __lowerCamelCase			:			Optional[Any]       			=					jax_cosine_distance(_a     ,      self.special_care_embeds  )
    __lowerCamelCase			:			Tuple       			=					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
    __lowerCamelCase			:			Dict       			=					0.0

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

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

    return has_nsfw_concepts






class    lowerCamelCase_       (							SCREAMING_SNAKE_CASE_	):
  """simple docstring"""




  a_				   =CLIPConfig
  a_				   ="""clip_input"""
  a_				   =FlaxStableDiffusionSafetyCheckerModule




  def __init__(					self :      List[str]     ,      _a :      CLIPConfig     ,      _a :      Optional[Tuple] = None     ,      _a :      int = 0     ,      _a :      jnp.dtype = jnp.floataa     ,      _a :      bool = True     ,      **_a :      List[Any]     ,      )   ->   List[Any]:
    if input_shape is None:
      __lowerCamelCase			:			str       			=					(1, 224, 224, 3)
    __lowerCamelCase			:			List[str]       			=					self.module_class(config=_a     ,      dtype=_a     ,      **_a  )
    super().__init__(_a     ,      _a     ,      input_shape=_a     ,      seed=_a     ,      dtype=_a     ,      _do_init=_do_init  )




  def 		_lowercase	(					self :      int     ,      _a :      jax.random.KeyArray     ,      _a :      Tuple     ,      _a :      FrozenDict = None  )   ->   FrozenDict:
    # init input tensor
    __lowerCamelCase			:			Optional[int]       			=					jax.random.normal(_a     ,      _a  )

    __lowerCamelCase			,__lowerCamelCase			:			Optional[Any]       			=					jax.random.split(_a  )
    __lowerCamelCase			:			int       			=					{'params': params_rng, 'dropout': dropout_rng}

    __lowerCamelCase			:			Tuple       			=					self.module.init(_a     ,      _a  )['params']

    return random_params




  def __call__(					self :      Optional[int]     ,      _a :      int     ,      _a :      dict = None     ,      )   ->   Dict:
    __lowerCamelCase			:			str       			=					jnp.transpose(_a     ,      (0, 2, 3, 1)  )

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






