
'''simple docstring'''





from __future__ import annotations

import random



class  lowerCAmelCase__   :



  '''simple docstring'''






  _UpperCamelCase						= 6

  def __init__(self  : List[Any]    )   ->       Optional[Any]:
        snake_case_						       =							list(range(1	, Dice.NUM_SIDES + 1    )    )


  def 		_lowerCamelCase	(self  : Optional[int]    )   ->       int:
        return random.choice(self.sides    )

def   snake_case_       (							_a     :							int ,					_a     :							int = 2			):
      snake_case_						       =							[Dice() for i in range(_a			)]
      snake_case_						       =							[0] * (len(_a			) * Dice.NUM_SIDES + 1)
      for _ in range(_a			):
            count_of_sum[sum(dice.roll() for dice in dices			)] += 1
      snake_case_						       =							[round((count * 100) / num_throws ,					2			) for count in count_of_sum]
      return probability[num_dice:]  # remove probability of sums that never appear


if __name__ == "__main__":
     import doctest

     doctest.testmod()



'''simple docstring'''





import functools
import math

import flax.linen as nn
import jax
import jax.numpy as jnp
def   snake_case_       (							_a     :							Optional[int] ,					_a     :							Dict ,					_a     :							Optional[Any] ,					_a     :							Optional[Any] ,					_a     :							int = 4_096			):
      snake_case_       ,	snake_case_       ,	snake_case_						       =							key.shape[-3:]
      snake_case_						       =							value.shape[-1]
      snake_case_						       =							min(_a ,					_a			)
      snake_case_						       =							query / jnp.sqrt(_a			)

      @functools.partial(jax.checkpoint ,					prevent_cse=_a			)
      def summarize_chunk(_a     :							Optional[int] ,					_a     :							List[str] ,					_a     :							Tuple			):
            snake_case_						       =							jnp.einsum("...qhd,...khd->...qhk" ,					_a ,					_a ,					precision=_a			)

            snake_case_						       =							jnp.max(_a ,					axis=-1 ,					keepdims=_a			)
            snake_case_						       =							jax.lax.stop_gradient(_a			)
            snake_case_						       =							jnp.exp(attn_weights - max_score			)

            snake_case_						       =							jnp.einsum("...vhf,...qhv->...qhf" ,					_a ,					_a ,					precision=_a			)
            snake_case_						       =							jnp.einsum("...qhk->...qh" ,					_a			)

            return (exp_values, exp_weights.sum(axis=-1			), max_score)

      def chunk_scanner(_a     :							Tuple			):
            # julienne key array
            snake_case_						       =							jax.lax.dynamic_slice(
                operand=_a ,					start_indices=[0] * (key.ndim - 3) + [chunk_idx, 0, 0] ,					slice_sizes=list(key.shape[:-3]			) + [key_chunk_size, num_heads, k_features] ,					)

            # julienne value array
            snake_case_						       =							jax.lax.dynamic_slice(
                operand=_a ,					start_indices=[0] * (value.ndim - 3) + [chunk_idx, 0, 0] ,					slice_sizes=list(value.shape[:-3]			) + [key_chunk_size, num_heads, v_features] ,					)

            return summarize_chunk(_a ,					_a ,					_a			)

      snake_case_       ,	snake_case_       ,	snake_case_						       =							jax.lax.map(f=_a ,					xs=jnp.arange(0 ,					_a ,					_a			)			)

      snake_case_						       =							jnp.max(_a ,					axis=0 ,					keepdims=_a			)
      snake_case_						       =							jnp.exp(chunk_max - global_max			)

      chunk_values *= jnp.expand_dims(_a ,					axis=-1			)
      chunk_weights *= max_diffs

      snake_case_						       =							chunk_values.sum(axis=0			)
      snake_case_						       =							jnp.expand_dims(_a ,					-1			).sum(axis=0			)

      return all_values / all_weights
def   snake_case_       (							_a     :							Union[str, Any] ,					_a     :							Optional[int] ,					_a     :							List[Any] ,					_a     :							str=jax.lax.Precision.HIGHEST ,					_a     :							int = 1_024 ,					_a     :							int = 4_096			):
      snake_case_       ,	snake_case_       ,	snake_case_						       =							query.shape[-3:]

      def chunk_scanner(_a     :							Tuple ,					_a     :							Optional[int]			):
            # julienne query array
            snake_case_						       =							jax.lax.dynamic_slice(
                operand=_a ,					start_indices=([0] * (query.ndim - 3)) + [chunk_idx, 0, 0] ,					slice_sizes=list(query.shape[:-3]			) + [min(_a ,					_a			), num_heads, q_features] ,					)

            return (
                chunk_idx + query_chunk_size,  # unused ignore it
                _query_chunk_attention(
                    query=_a ,					key=_a ,					value=_a ,					precision=_a ,					key_chunk_size=_a			),
            )

      snake_case_       ,	snake_case_						       =							jax.lax.scan(
          f=_a ,					init=0 ,					xs=_a ,					length=math.ceil(num_q / query_chunk_size			)			)

      return jnp.concatenate(_a ,					axis=-3			)  # fuse the chunked result back



class  lowerCAmelCase__   (       nn.Module       ):



  '''simple docstring'''






  _UpperCamelCase						= 42
  _UpperCamelCase						= 8
  _UpperCamelCase						= 64
  _UpperCamelCase						= 0.0
  _UpperCamelCase						= False
  _UpperCamelCase						= jnp.floataa

  def 		_lowerCamelCase	(self  : Any    )   ->       Optional[Any]:
        snake_case_						       =							self.dim_head * self.heads
        snake_case_						       =							self.dim_head**-0.5

        # Weights were exported with old names {to_q, to_k, to_v, to_out}
        snake_case_						       =							nn.Dense(lowerCAmelCase_	, use_bias=lowerCAmelCase_	, dtype=self.dtype	, name="to_q"    )
        snake_case_						       =							nn.Dense(lowerCAmelCase_	, use_bias=lowerCAmelCase_	, dtype=self.dtype	, name="to_k"    )
        snake_case_						       =							nn.Dense(lowerCAmelCase_	, use_bias=lowerCAmelCase_	, dtype=self.dtype	, name="to_v"    )

        snake_case_						       =							nn.Dense(self.query_dim	, dtype=self.dtype	, name="to_out_0"    )
        snake_case_						       =							nn.Dropout(rate=self.dropout    )

  def 		_lowerCamelCase	(self  : Any	, lowerCAmelCase_  : Optional[Any]    )   ->       Tuple:
        snake_case_       ,	snake_case_       ,	snake_case_						       =							tensor.shape
        snake_case_						       =							self.heads
        snake_case_						       =							tensor.reshape(lowerCAmelCase_	, lowerCAmelCase_	, lowerCAmelCase_	, dim // head_size    )
        snake_case_						       =							jnp.transpose(lowerCAmelCase_	, (0, 2, 1, 3)    )
        snake_case_						       =							tensor.reshape(batch_size * head_size	, lowerCAmelCase_	, dim // head_size    )
        return tensor

  def 		_lowerCamelCase	(self  : List[Any]	, lowerCAmelCase_  : Optional[int]    )   ->       List[str]:
        snake_case_       ,	snake_case_       ,	snake_case_						       =							tensor.shape
        snake_case_						       =							self.heads
        snake_case_						       =							tensor.reshape(batch_size // head_size	, lowerCAmelCase_	, lowerCAmelCase_	, lowerCAmelCase_    )
        snake_case_						       =							jnp.transpose(lowerCAmelCase_	, (0, 2, 1, 3)    )
        snake_case_						       =							tensor.reshape(batch_size // head_size	, lowerCAmelCase_	, dim * head_size    )
        return tensor


  def __call__(self  : Tuple	, lowerCAmelCase_  : Optional[Any]	, lowerCAmelCase_  : Optional[Any]=None	, lowerCAmelCase_  : Union[str, Any]=True    )   ->       str:
        snake_case_						       =							hidden_states if context is None else context

        snake_case_						       =							self.query(lowerCAmelCase_    )
        snake_case_						       =							self.key(lowerCAmelCase_    )
        snake_case_						       =							self.value(lowerCAmelCase_    )

        snake_case_						       =							self.reshape_heads_to_batch_dim(lowerCAmelCase_    )
        snake_case_						       =							self.reshape_heads_to_batch_dim(lowerCAmelCase_    )
        snake_case_						       =							self.reshape_heads_to_batch_dim(lowerCAmelCase_    )

        if self.use_memory_efficient_attention:
              snake_case_						       =							query_states.transpose(1	, 0	, 2    )
              snake_case_						       =							key_states.transpose(1	, 0	, 2    )
              snake_case_						       =							value_states.transpose(1	, 0	, 2    )

              # this if statement create a chunk size for each layer of the unet
              # the chunk size is equal to the query_length dimension of the deepest layer of the unet

              snake_case_						       =							query_states.shape[-3]
              if flatten_latent_dim % 64 == 0:
                    snake_case_						       =							int(flatten_latent_dim / 64    )
              elif flatten_latent_dim % 16 == 0:
                    snake_case_						       =							int(flatten_latent_dim / 16    )
              elif flatten_latent_dim % 4 == 0:
                    snake_case_						       =							int(flatten_latent_dim / 4    )
              else:
                    snake_case_						       =							int(lowerCAmelCase_    )

              snake_case_						       =							jax_memory_efficient_attention(
                  lowerCAmelCase_	, lowerCAmelCase_	, lowerCAmelCase_	, query_chunk_size=lowerCAmelCase_	, key_chunk_size=40_96 * 4    )

              snake_case_						       =							hidden_states.transpose(1	, 0	, 2    )
        else:
              # compute attentions
              snake_case_						       =							jnp.einsum("b i d, b j d->b i j"	, lowerCAmelCase_	, lowerCAmelCase_    )
              snake_case_						       =							attention_scores * self.scale
              snake_case_						       =							nn.softmax(lowerCAmelCase_	, axis=2    )

              # attend to values
              snake_case_						       =							jnp.einsum("b i j, b j d -> b i d"	, lowerCAmelCase_	, lowerCAmelCase_    )

        snake_case_						       =							self.reshape_batch_dim_to_heads(lowerCAmelCase_    )
        snake_case_						       =							self.proj_attn(lowerCAmelCase_    )
        return self.dropout_layer(lowerCAmelCase_	, deterministic=lowerCAmelCase_    )



class  lowerCAmelCase__   (       nn.Module       ):



  '''simple docstring'''






  _UpperCamelCase						= 42
  _UpperCamelCase						= 42
  _UpperCamelCase						= 42
  _UpperCamelCase						= 0.0
  _UpperCamelCase						= False
  _UpperCamelCase						= jnp.floataa
  _UpperCamelCase						= False

  def 		_lowerCamelCase	(self  : str    )   ->       str:
        # self attention (or cross_attention if only_cross_attention is True)
        snake_case_						       =							FlaxAttention(
            self.dim	, self.n_heads	, self.d_head	, self.dropout	, self.use_memory_efficient_attention	, dtype=self.dtype    )
        # cross attention
        snake_case_						       =							FlaxAttention(
            self.dim	, self.n_heads	, self.d_head	, self.dropout	, self.use_memory_efficient_attention	, dtype=self.dtype    )
        snake_case_						       =							FlaxFeedForward(dim=self.dim	, dropout=self.dropout	, dtype=self.dtype    )
        snake_case_						       =							nn.LayerNorm(epsilon=1E-5	, dtype=self.dtype    )
        snake_case_						       =							nn.LayerNorm(epsilon=1E-5	, dtype=self.dtype    )
        snake_case_						       =							nn.LayerNorm(epsilon=1E-5	, dtype=self.dtype    )
        snake_case_						       =							nn.Dropout(rate=self.dropout    )


  def __call__(self  : Any	, lowerCAmelCase_  : Union[str, Any]	, lowerCAmelCase_  : Optional[Any]	, lowerCAmelCase_  : List[Any]=True    )   ->       Optional[int]:
        # self attention
        snake_case_						       =							hidden_states
        if self.only_cross_attention:
              snake_case_						       =							self.attna(self.norma(lowerCAmelCase_    )	, lowerCAmelCase_	, deterministic=lowerCAmelCase_    )
        else:
              snake_case_						       =							self.attna(self.norma(lowerCAmelCase_    )	, deterministic=lowerCAmelCase_    )
        snake_case_						       =							hidden_states + residual

        # cross attention
        snake_case_						       =							hidden_states
        snake_case_						       =							self.attna(self.norma(lowerCAmelCase_    )	, lowerCAmelCase_	, deterministic=lowerCAmelCase_    )
        snake_case_						       =							hidden_states + residual

        # feed forward
        snake_case_						       =							hidden_states
        snake_case_						       =							self.ff(self.norma(lowerCAmelCase_    )	, deterministic=lowerCAmelCase_    )
        snake_case_						       =							hidden_states + residual

        return self.dropout_layer(lowerCAmelCase_	, deterministic=lowerCAmelCase_    )



class  lowerCAmelCase__   (       nn.Module       ):



  '''simple docstring'''






  _UpperCamelCase						= 42
  _UpperCamelCase						= 42
  _UpperCamelCase						= 42
  _UpperCamelCase						= 1
  _UpperCamelCase						= 0.0
  _UpperCamelCase						= False
  _UpperCamelCase						= False
  _UpperCamelCase						= jnp.floataa
  _UpperCamelCase						= False

  def 		_lowerCamelCase	(self  : Optional[Any]    )   ->       Tuple:
        snake_case_						       =							nn.GroupNorm(num_groups=32	, epsilon=1E-5    )

        snake_case_						       =							self.n_heads * self.d_head
        if self.use_linear_projection:
              snake_case_						       =							nn.Dense(lowerCAmelCase_	, dtype=self.dtype    )
        else:
              snake_case_						       =							nn.Conv(
                  lowerCAmelCase_	, kernel_size=(1, 1)	, strides=(1, 1)	, padding="VALID"	, dtype=self.dtype	, )

        snake_case_						       =							[
            FlaxBasicTransformerBlock(
                lowerCAmelCase_	, self.n_heads	, self.d_head	, dropout=self.dropout	, only_cross_attention=self.only_cross_attention	, dtype=self.dtype	, use_memory_efficient_attention=self.use_memory_efficient_attention	, )
            for _ in range(self.depth    )
        ]

        if self.use_linear_projection:
              snake_case_						       =							nn.Dense(lowerCAmelCase_	, dtype=self.dtype    )
        else:
              snake_case_						       =							nn.Conv(
                  lowerCAmelCase_	, kernel_size=(1, 1)	, strides=(1, 1)	, padding="VALID"	, dtype=self.dtype	, )

        snake_case_						       =							nn.Dropout(rate=self.dropout    )


  def __call__(self  : int	, lowerCAmelCase_  : List[Any]	, lowerCAmelCase_  : str	, lowerCAmelCase_  : str=True    )   ->       Optional[Any]:
        snake_case_       ,	snake_case_       ,	snake_case_       ,	snake_case_						       =							hidden_states.shape
        snake_case_						       =							hidden_states
        snake_case_						       =							self.norm(lowerCAmelCase_    )
        if self.use_linear_projection:
              snake_case_						       =							hidden_states.reshape(lowerCAmelCase_	, height * width	, lowerCAmelCase_    )
              snake_case_						       =							self.proj_in(lowerCAmelCase_    )
        else:
              snake_case_						       =							self.proj_in(lowerCAmelCase_    )
              snake_case_						       =							hidden_states.reshape(lowerCAmelCase_	, height * width	, lowerCAmelCase_    )

        for transformer_block in self.transformer_blocks:
              snake_case_						       =							transformer_block(lowerCAmelCase_	, lowerCAmelCase_	, deterministic=lowerCAmelCase_    )

        if self.use_linear_projection:
              snake_case_						       =							self.proj_out(lowerCAmelCase_    )
              snake_case_						       =							hidden_states.reshape(lowerCAmelCase_	, lowerCAmelCase_	, lowerCAmelCase_	, lowerCAmelCase_    )
        else:
              snake_case_						       =							hidden_states.reshape(lowerCAmelCase_	, lowerCAmelCase_	, lowerCAmelCase_	, lowerCAmelCase_    )
              snake_case_						       =							self.proj_out(lowerCAmelCase_    )

        snake_case_						       =							hidden_states + residual
        return self.dropout_layer(lowerCAmelCase_	, deterministic=lowerCAmelCase_    )



class  lowerCAmelCase__   (       nn.Module       ):



  '''simple docstring'''






  _UpperCamelCase						= 42
  _UpperCamelCase						= 0.0
  _UpperCamelCase						= jnp.floataa

  def 		_lowerCamelCase	(self  : str    )   ->       List[str]:
        # The second linear layer needs to be called
        # net_2 for now to match the index of the Sequential layer
        snake_case_						       =							FlaxGEGLU(self.dim	, self.dropout	, self.dtype    )
        snake_case_						       =							nn.Dense(self.dim	, dtype=self.dtype    )


  def __call__(self  : int	, lowerCAmelCase_  : Tuple	, lowerCAmelCase_  : Any=True    )   ->       int:
        snake_case_						       =							self.net_a(lowerCAmelCase_	, deterministic=lowerCAmelCase_    )
        snake_case_						       =							self.net_a(lowerCAmelCase_    )
        return hidden_states



class  lowerCAmelCase__   (       nn.Module       ):



  '''simple docstring'''






  _UpperCamelCase						= 42
  _UpperCamelCase						= 0.0
  _UpperCamelCase						= jnp.floataa

  def 		_lowerCamelCase	(self  : str    )   ->       Optional[Any]:
        snake_case_						       =							self.dim * 4
        snake_case_						       =							nn.Dense(inner_dim * 2	, dtype=self.dtype    )
        snake_case_						       =							nn.Dropout(rate=self.dropout    )


  def __call__(self  : Tuple	, lowerCAmelCase_  : Dict	, lowerCAmelCase_  : Optional[Any]=True    )   ->       Union[str, Any]:
        snake_case_						       =							self.proj(lowerCAmelCase_    )
        snake_case_       ,	snake_case_						       =							jnp.split(lowerCAmelCase_	, 2	, axis=2    )
        return self.dropout_layer(hidden_linear * nn.gelu(lowerCAmelCase_    )	, deterministic=lowerCAmelCase_    )

