

import math


def     _UpperCAmelCase				(     _SCREAMING_SNAKE_CASE    :						int				):





    """simple docstring"""


    if 1 < number < 4:
        # 2 and 3 are primes
        return True
    elif number < 2 or number % 2 == 0 or number % 3 == 0:
        # Negatives, 0, 1, all even numbers, all multiples of 3 are not primes
        return False

    # All primes number are in format of 6k +/- 1
    for i in range(5       ,      int(math.sqrt(_SCREAMING_SNAKE_CASE				) + 1				)       ,      6				):
        if number % i == 0 or number % (i + 2) == 0:
            return False
    return True


def     _UpperCAmelCase				(     _SCREAMING_SNAKE_CASE    :						int = 600_851_475_143				):





    """simple docstring"""


    try:
        SCREAMING_SNAKE_CASE_													=  int(_SCREAMING_SNAKE_CASE				)
    except (TypeError, ValueError):
        raise TypeError('Parameter n must be int or castable to int.'				)
    if n <= 0:
        raise ValueError('Parameter n must be greater than or equal to one.'				)
    SCREAMING_SNAKE_CASE_													=  0
    if is_prime(_SCREAMING_SNAKE_CASE				):
        return n
    while n % 2 == 0:
        n //= 2
    if is_prime(_SCREAMING_SNAKE_CASE				):
        return n
    for i in range(3       ,      int(math.sqrt(_SCREAMING_SNAKE_CASE				)				) + 1       ,      2				):
        if n % i == 0:
            if is_prime(n // i				):
                SCREAMING_SNAKE_CASE_													=  n // i
                break
            elif is_prime(_SCREAMING_SNAKE_CASE				):
                SCREAMING_SNAKE_CASE_													=  i
    return max_number


if __name__ == "__main__":
       print(F'{solution() = }')





from dataclasses import dataclass
from typing import Any, Dict, Optional

import torch
import torch.nn.functional as F
from torch import nn

from ..configuration_utils import ConfigMixin, register_to_config
from ..models.embeddings import ImagePositionalEmbeddings
from ..utils import BaseOutput, deprecate
from .attention import BasicTransformerBlock
from .embeddings import PatchEmbed
from .modeling_utils import ModelMixin



@dataclass
class 					__snake_case     (			lowerCAmelCase__				):
  __lowerCAmelCase   :		torch.FloatTensor






class 					__snake_case     (			lowerCAmelCase__   ,       lowerCAmelCase__				):

  @register_to_config
  def __init__(    self  ,						_A = 16  ,						_A = 88  ,						_A = None  ,						_A = None  ,						_A = 1  ,						_A = 0.0  ,						_A = 32  ,						_A = None  ,						_A = False  ,						_A = None  ,						_A = None  ,						_A = None  ,						_A = "geglu"  ,						_A = None  ,						_A = False  ,						_A = False  ,						_A = False  ,						_A = "layer_norm"  ,						_A = True  ,						):
      super().__init__()
      SCREAMING_SNAKE_CASE_													=  use_linear_projection
      SCREAMING_SNAKE_CASE_													=  num_attention_heads
      SCREAMING_SNAKE_CASE_													=  attention_head_dim
      SCREAMING_SNAKE_CASE_													=  num_attention_heads * attention_head_dim

      # 1. Transformer2DModel can process both standard continuous images of shape `(batch_size, num_channels, width, height)` as well as quantized image embeddings of shape `(batch_size, num_image_vectors)`
      # Define whether input is continuous or discrete depending on configuration
      SCREAMING_SNAKE_CASE_													=  (in_channels is not None) and (patch_size is None)
      SCREAMING_SNAKE_CASE_													=  num_vector_embeds is not None
      SCREAMING_SNAKE_CASE_													=  in_channels is not None and patch_size is not None

      if norm_type == "layer_norm" and num_embeds_ada_norm is not None:
          SCREAMING_SNAKE_CASE_													=  (
              f"""The configuration file of this model: {self.__class__} is outdated. `norm_type` is either not set or"""
              ' incorrectly set to `\'layer_norm\'`.Make sure to set `norm_type` to `\'ada_norm\'` in the config.'
              ' Please make sure to update the config accordingly as leaving `norm_type` might led to incorrect'
              ' results in future versions. If you have downloaded this checkpoint from the Hugging Face Hub, it'
              ' would be very nice if you could open a Pull request for the `transformer/config.json` file'
          )
          deprecate('norm_type!=num_embeds_ada_norm'  ,						'1.0.0'  ,						_A  ,						standard_warn=_A)
          SCREAMING_SNAKE_CASE_													=  'ada_norm'

      if self.is_input_continuous and self.is_input_vectorized:
          raise ValueError(
              f"""Cannot define both `in_channels`: {in_channels} and `num_vector_embeds`: {num_vector_embeds}. Make"""
              ' sure that either `in_channels` or `num_vector_embeds` is None.')
      elif self.is_input_vectorized and self.is_input_patches:
          raise ValueError(
              f"""Cannot define both `num_vector_embeds`: {num_vector_embeds} and `patch_size`: {patch_size}. Make"""
              ' sure that either `num_vector_embeds` or `num_patches` is None.')
      elif not self.is_input_continuous and not self.is_input_vectorized and not self.is_input_patches:
          raise ValueError(
              f"""Has to define `in_channels`: {in_channels}, `num_vector_embeds`: {num_vector_embeds}, or patch_size:"""
              f""" {patch_size}. Make sure that `in_channels`, `num_vector_embeds` or `num_patches` is not None.""")

      # 2. Define input layers
      if self.is_input_continuous:
          SCREAMING_SNAKE_CASE_													=  in_channels

          SCREAMING_SNAKE_CASE_													=  torch.nn.GroupNorm(num_groups=_A  ,						num_channels=_A  ,						eps=1E-6  ,						affine=_A)
          if use_linear_projection:
              SCREAMING_SNAKE_CASE_													=  nn.Linear(_A  ,						_A)
          else:
              SCREAMING_SNAKE_CASE_													=  nn.Convad(_A  ,						_A  ,						kernel_size=1  ,						stride=1  ,						padding=0)
      elif self.is_input_vectorized:
          assert sample_size is not None, "Transformer2DModel over discrete input must provide sample_size"
          assert num_vector_embeds is not None, "Transformer2DModel over discrete input must provide num_embed"

          SCREAMING_SNAKE_CASE_													=  sample_size
          SCREAMING_SNAKE_CASE_													=  sample_size
          SCREAMING_SNAKE_CASE_													=  num_vector_embeds
          SCREAMING_SNAKE_CASE_													=  self.height * self.width

          SCREAMING_SNAKE_CASE_													=  ImagePositionalEmbeddings(
              num_embed=_A  ,						embed_dim=_A  ,						height=self.height  ,						width=self.width)
      elif self.is_input_patches:
          assert sample_size is not None, "Transformer2DModel over patched input must provide sample_size"

          SCREAMING_SNAKE_CASE_													=  sample_size
          SCREAMING_SNAKE_CASE_													=  sample_size

          SCREAMING_SNAKE_CASE_													=  patch_size
          SCREAMING_SNAKE_CASE_													=  PatchEmbed(
              height=_A  ,						width=_A  ,						patch_size=_A  ,						in_channels=_A  ,						embed_dim=_A  ,						)

      # 3. Define transformers blocks
      SCREAMING_SNAKE_CASE_													=  nn.ModuleList(
          [
              BasicTransformerBlock(
                  _A  ,						_A  ,						_A  ,						dropout=_A  ,						cross_attention_dim=_A  ,						activation_fn=_A  ,						num_embeds_ada_norm=_A  ,						attention_bias=_A  ,						only_cross_attention=_A  ,						upcast_attention=_A  ,						norm_type=_A  ,						norm_elementwise_affine=_A  ,						)
              for d in range(_A)
          ])

      # 4. Define output layers
      SCREAMING_SNAKE_CASE_													=  in_channels if out_channels is None else out_channels
      if self.is_input_continuous:
          # TODO: should use out_channels for continuous projections
          if use_linear_projection:
              SCREAMING_SNAKE_CASE_													=  nn.Linear(_A  ,						_A)
          else:
              SCREAMING_SNAKE_CASE_													=  nn.Convad(_A  ,						_A  ,						kernel_size=1  ,						stride=1  ,						padding=0)
      elif self.is_input_vectorized:
          SCREAMING_SNAKE_CASE_													=  nn.LayerNorm(_A)
          SCREAMING_SNAKE_CASE_													=  nn.Linear(_A  ,						self.num_vector_embeds - 1)
      elif self.is_input_patches:
          SCREAMING_SNAKE_CASE_													=  nn.LayerNorm(_A  ,						elementwise_affine=_A  ,						eps=1E-6)
          SCREAMING_SNAKE_CASE_													=  nn.Linear(_A  ,						2 * inner_dim)
          SCREAMING_SNAKE_CASE_													=  nn.Linear(_A  ,						patch_size * patch_size * self.out_channels)


  def   lowerCAmelCase__   (    self  ,						_A  ,						_A = None  ,						_A = None  ,						_A = None  ,						_A = None  ,						_A = None  ,						_A = None  ,						_A = True  ,						):
      # ensure attention_mask is a bias, and give it a singleton query_tokens dimension.
      #   we may have done this conversion already, e.g. if we came here via UNet2DConditionModel#forward.
      #   we can tell by counting dims; if ndim == 2: it's a mask rather than a bias.
      # expects mask of shape:
      #   [batch, key_tokens]
      # adds singleton query_tokens dimension:
      #   [batch,                    1, key_tokens]
      # this helps to broadcast it as a bias over attention scores, which will be in one of the following shapes:
      #   [batch,  heads, query_tokens, key_tokens] (e.g. torch sdp attn)
      #   [batch * heads, query_tokens, key_tokens] (e.g. xformers or classic attn)
      if attention_mask is not None and attention_mask.ndim == 2:
          # assume that mask is expressed as:
          #   (1 = keep,      0 = discard)
          # convert mask into a bias that can be added to attention scores:
          #       (keep = +0,     discard = -10000.0)
          SCREAMING_SNAKE_CASE_													=  (1 - attention_mask.to(hidden_states.dtype)) * -1_0_0_0_0.0
          SCREAMING_SNAKE_CASE_													=  attention_mask.unsqueeze(1)

      # convert encoder_attention_mask to a bias the same way we do for attention_mask
      if encoder_attention_mask is not None and encoder_attention_mask.ndim == 2:
          SCREAMING_SNAKE_CASE_													=  (1 - encoder_attention_mask.to(hidden_states.dtype)) * -1_0_0_0_0.0
          SCREAMING_SNAKE_CASE_													=  encoder_attention_mask.unsqueeze(1)

      # 1. Input
      if self.is_input_continuous:
          SCREAMING_SNAKE_CASE_							,	SCREAMING_SNAKE_CASE_							,	SCREAMING_SNAKE_CASE_							,	SCREAMING_SNAKE_CASE_													=  hidden_states.shape
          SCREAMING_SNAKE_CASE_													=  hidden_states

          SCREAMING_SNAKE_CASE_													=  self.norm(_A)
          if not self.use_linear_projection:
              SCREAMING_SNAKE_CASE_													=  self.proj_in(_A)
              SCREAMING_SNAKE_CASE_													=  hidden_states.shape[1]
              SCREAMING_SNAKE_CASE_													=  hidden_states.permute(0  ,						2  ,						3  ,						1).reshape(_A  ,						height * width  ,						_A)
          else:
              SCREAMING_SNAKE_CASE_													=  hidden_states.shape[1]
              SCREAMING_SNAKE_CASE_													=  hidden_states.permute(0  ,						2  ,						3  ,						1).reshape(_A  ,						height * width  ,						_A)
              SCREAMING_SNAKE_CASE_													=  self.proj_in(_A)
      elif self.is_input_vectorized:
          SCREAMING_SNAKE_CASE_													=  self.latent_image_embedding(_A)
      elif self.is_input_patches:
          SCREAMING_SNAKE_CASE_													=  self.pos_embed(_A)

      # 2. Blocks
      for block in self.transformer_blocks:
          SCREAMING_SNAKE_CASE_													=  block(
              _A  ,						attention_mask=_A  ,						encoder_hidden_states=_A  ,						encoder_attention_mask=_A  ,						timestep=_A  ,						cross_attention_kwargs=_A  ,						class_labels=_A  ,						)

      # 3. Output
      if self.is_input_continuous:
          if not self.use_linear_projection:
              SCREAMING_SNAKE_CASE_													=  hidden_states.reshape(_A  ,						_A  ,						_A  ,						_A).permute(0  ,						3  ,						1  ,						2).contiguous()
              SCREAMING_SNAKE_CASE_													=  self.proj_out(_A)
          else:
              SCREAMING_SNAKE_CASE_													=  self.proj_out(_A)
              SCREAMING_SNAKE_CASE_													=  hidden_states.reshape(_A  ,						_A  ,						_A  ,						_A).permute(0  ,						3  ,						1  ,						2).contiguous()

          SCREAMING_SNAKE_CASE_													=  hidden_states + residual
      elif self.is_input_vectorized:
          SCREAMING_SNAKE_CASE_													=  self.norm_out(_A)
          SCREAMING_SNAKE_CASE_													=  self.out(_A)
          # (batch, self.num_vector_embeds - 1, self.num_latent_pixels)
          SCREAMING_SNAKE_CASE_													=  logits.permute(0  ,						2  ,						1)

          # log(p(x_0))
          SCREAMING_SNAKE_CASE_													=  F.log_softmax(logits.double()  ,						dim=1).float()
      elif self.is_input_patches:
          # TODO: cleanup!
          SCREAMING_SNAKE_CASE_													=  self.transformer_blocks[0].norma.emb(
              _A  ,						_A  ,						hidden_dtype=hidden_states.dtype)
          SCREAMING_SNAKE_CASE_							,	SCREAMING_SNAKE_CASE_													=  self.proj_out_a(F.silu(_A)).chunk(2  ,						dim=1)
          SCREAMING_SNAKE_CASE_													=  self.norm_out(_A) * (1 + scale[:, None]) + shift[:, None]
          SCREAMING_SNAKE_CASE_													=  self.proj_out_a(_A)

          # unpatchify
          SCREAMING_SNAKE_CASE_													=  SCREAMING_SNAKE_CASE_													=  int(hidden_states.shape[1] ** 0.5)
          SCREAMING_SNAKE_CASE_													=  hidden_states.reshape(
              shape=(-1, height, width, self.patch_size, self.patch_size, self.out_channels))
          SCREAMING_SNAKE_CASE_													=  torch.einsum('nhwpqc->nchpwq'  ,						_A)
          SCREAMING_SNAKE_CASE_													=  hidden_states.reshape(
              shape=(-1, self.out_channels, height * self.patch_size, width * self.patch_size))

      if not return_dict:
          return (output,)

      return TransformeraDModelOutput(sample=_A)


