code: 
  
  import functools
  import math
  
  import flax.linen as nn
  import jax
  import jax.numpy as jnp
  
  
  
  def  _A  (			SCREAMING_SNAKE_CASE  ,SCREAMING_SNAKE_CASE  ,SCREAMING_SNAKE_CASE  ,SCREAMING_SNAKE_CASE  ,SCREAMING_SNAKE_CASE = 4_0_9_6			):
  	UpperCAmelCase__       ,	UpperCAmelCase__       ,	UpperCAmelCase__:  List[Any]		       =   key.shape[-3:]
  	UpperCAmelCase__:  Optional[int]		       =   value.shape[-1]
  	UpperCAmelCase__:  Dict		       =   min(SCREAMING_SNAKE_CASE  ,SCREAMING_SNAKE_CASE			)
  	UpperCAmelCase__:  List[Any]		       =   query / jnp.sqrt(SCREAMING_SNAKE_CASE			)
  
  	@functools.partial(jax.checkpoint  ,prevent_cse=SCREAMING_SNAKE_CASE			)
  	def summarize_chunk(SCREAMING_SNAKE_CASE  ,SCREAMING_SNAKE_CASE  ,SCREAMING_SNAKE_CASE			):
  		UpperCAmelCase__:  Tuple		       =   jnp.einsum("...qhd,...khd->...qhk"  ,SCREAMING_SNAKE_CASE  ,SCREAMING_SNAKE_CASE  ,precision=SCREAMING_SNAKE_CASE			)
  
  		UpperCAmelCase__:  Any		       =   jnp.max(SCREAMING_SNAKE_CASE  ,axis=-1  ,keepdims=SCREAMING_SNAKE_CASE			)
  		UpperCAmelCase__:  str		       =   jax.lax.stop_gradient(SCREAMING_SNAKE_CASE			)
  		UpperCAmelCase__:  Optional[int]		       =   jnp.exp(attn_weights - max_score			)
  
  		UpperCAmelCase__:  str		       =   jnp.einsum("...vhf,...qhv->...qhf"  ,SCREAMING_SNAKE_CASE  ,SCREAMING_SNAKE_CASE  ,precision=SCREAMING_SNAKE_CASE			)
  		UpperCAmelCase__:  Tuple		       =   jnp.einsum("...qhk->...qh"  ,SCREAMING_SNAKE_CASE			)
  
  		return (exp_values, exp_weights.sum(axis=-1			), max_score)
  
  	def chunk_scanner(SCREAMING_SNAKE_CASE			):
  		# julienne key array
  		UpperCAmelCase__:  str		       =   jax.lax.dynamic_slice(
  		    operand=SCREAMING_SNAKE_CASE  ,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
  		UpperCAmelCase__:  List[Any]		       =   jax.lax.dynamic_slice(
  		    operand=SCREAMING_SNAKE_CASE  ,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(SCREAMING_SNAKE_CASE  ,SCREAMING_SNAKE_CASE  ,SCREAMING_SNAKE_CASE			)
  
  	UpperCAmelCase__       ,	UpperCAmelCase__       ,	UpperCAmelCase__:  Optional[int]		       =   jax.lax.map(f=SCREAMING_SNAKE_CASE  ,xs=jnp.arange(0  ,SCREAMING_SNAKE_CASE  ,SCREAMING_SNAKE_CASE			)			)
  
  	UpperCAmelCase__:  Optional[Any]		       =   jnp.max(SCREAMING_SNAKE_CASE  ,axis=0  ,keepdims=SCREAMING_SNAKE_CASE			)
  	UpperCAmelCase__:  Union[str, Any]		       =   jnp.exp(chunk_max - global_max			)
  
  	chunk_values *= jnp.expand_dims(SCREAMING_SNAKE_CASE  ,axis=-1			)
  	chunk_weights *= max_diffs
  
  	UpperCAmelCase__:  Optional[Any]		       =   chunk_values.sum(axis=0			)
  	UpperCAmelCase__:  Tuple		       =   jnp.expand_dims(SCREAMING_SNAKE_CASE  ,-1			).sum(axis=0			)
  
  	return all_values / all_weights
  
  
  
  def  _A  (			SCREAMING_SNAKE_CASE  ,SCREAMING_SNAKE_CASE  ,SCREAMING_SNAKE_CASE  ,SCREAMING_SNAKE_CASE=jax.lax.Precision.HIGHEST  ,SCREAMING_SNAKE_CASE = 1_0_2_4  ,SCREAMING_SNAKE_CASE = 4_0_9_6			):
  	UpperCAmelCase__       ,	UpperCAmelCase__       ,	UpperCAmelCase__:  Optional[int]		       =   query.shape[-3:]
  
  	def chunk_scanner(SCREAMING_SNAKE_CASE  ,SCREAMING_SNAKE_CASE			):
  		# julienne query array
  		UpperCAmelCase__:  List[str]		       =   jax.lax.dynamic_slice(
  		    operand=SCREAMING_SNAKE_CASE  ,start_indices=([0] * (query.ndim - 3)) + [chunk_idx, 0, 0]  ,slice_sizes=list(query.shape[:-3]			) + [min(SCREAMING_SNAKE_CASE  ,SCREAMING_SNAKE_CASE			), num_heads, q_features]  ,)
  
  		return (
  		    chunk_idx + query_chunk_size,  # unused ignore it
  		    _query_chunk_attention(
  		        query=SCREAMING_SNAKE_CASE  ,key=SCREAMING_SNAKE_CASE  ,value=SCREAMING_SNAKE_CASE  ,precision=SCREAMING_SNAKE_CASE  ,key_chunk_size=SCREAMING_SNAKE_CASE			),
  		)
  
  	UpperCAmelCase__       ,	UpperCAmelCase__:  List[str]		       =   jax.lax.scan(
  	    f=SCREAMING_SNAKE_CASE  ,init=0  ,xs=SCREAMING_SNAKE_CASE  ,length=math.ceil(num_q / query_chunk_size			)			)
  
  	return jnp.concatenate(SCREAMING_SNAKE_CASE  ,axis=-3			)  # fuse the chunked result back
  
  class __UpperCamelCase    (							nn.Module		):
  
  
  
  
  
  
  					'''simple docstring'''
  
  
  
  					__magic_name__     			=      42
  					__magic_name__     			=      8
  					__magic_name__     			=      6_4
  					__magic_name__     			=      0.0
  					__magic_name__     			=      False
  					__magic_name__     			=      jnp.floataa
  
  
  
  
  
  					def       _UpperCAmelCase						(							self	):
  						UpperCAmelCase__:  Optional[Any]		       =   self.dim_head * self.heads
  						UpperCAmelCase__:  Dict		       =   self.dim_head**-0.5
  
  						# Weights were exported with old names {to_q, to_k, to_v, to_out}
  						UpperCAmelCase__:  str		       =   nn.Dense(lowerCamelCase__							,			use_bias=lowerCamelCase__							,			dtype=self.dtype							,			name="to_q"	)
  						UpperCAmelCase__:  Union[str, Any]		       =   nn.Dense(lowerCamelCase__							,			use_bias=lowerCamelCase__							,			dtype=self.dtype							,			name="to_k"	)
  						UpperCAmelCase__:  List[Any]		       =   nn.Dense(lowerCamelCase__							,			use_bias=lowerCamelCase__							,			dtype=self.dtype							,			name="to_v"	)
  
  						UpperCAmelCase__:  int		       =   nn.Dense(self.query_dim							,			dtype=self.dtype							,			name="to_out_0"	)
  						UpperCAmelCase__:  str		       =   nn.Dropout(rate=self.dropout	)
  
  
  
  
  
  					def       _UpperCAmelCase						(							self							,			lowerCamelCase__	):
  						UpperCAmelCase__       ,	UpperCAmelCase__       ,	UpperCAmelCase__:  Any		       =   tensor.shape
  						UpperCAmelCase__:  List[Any]		       =   self.heads
  						UpperCAmelCase__:  Any		       =   tensor.reshape(lowerCamelCase__							,			lowerCamelCase__							,			lowerCamelCase__							,			dim // head_size	)
  						UpperCAmelCase__:  Union[str, Any]		       =   jnp.transpose(lowerCamelCase__							,			(0, 2, 1, 3)	)
  						UpperCAmelCase__:  Dict		       =   tensor.reshape(batch_size * head_size							,			lowerCamelCase__							,			dim // head_size	)
  						return tensor
  
  
  
  
  
  					def       _UpperCAmelCase						(							self							,			lowerCamelCase__	):
  						UpperCAmelCase__       ,	UpperCAmelCase__       ,	UpperCAmelCase__:  Union[str, Any]		       =   tensor.shape
  						UpperCAmelCase__:  int		       =   self.heads
  						UpperCAmelCase__:  Union[str, Any]		       =   tensor.reshape(batch_size // head_size							,			lowerCamelCase__							,			lowerCamelCase__							,			lowerCamelCase__	)
  						UpperCAmelCase__:  Tuple		       =   jnp.transpose(lowerCamelCase__							,			(0, 2, 1, 3)	)
  						UpperCAmelCase__:  List[Any]		       =   tensor.reshape(batch_size // head_size							,			lowerCamelCase__							,			dim * head_size	)
  						return tensor
  
  
  
  
  
  					def __call__(							self							,			lowerCamelCase__							,			lowerCamelCase__=None							,			lowerCamelCase__=True	):
  						UpperCAmelCase__:  str		       =   hidden_states if context is None else context
  
  						UpperCAmelCase__:  Union[str, Any]		       =   self.query(lowerCamelCase__	)
  						UpperCAmelCase__:  Optional[int]		       =   self.key(lowerCamelCase__	)
  						UpperCAmelCase__:  Optional[int]		       =   self.value(lowerCamelCase__	)
  
  						UpperCAmelCase__:  str		       =   self.reshape_heads_to_batch_dim(lowerCamelCase__	)
  						UpperCAmelCase__:  Optional[int]		       =   self.reshape_heads_to_batch_dim(lowerCamelCase__	)
  						UpperCAmelCase__:  List[str]		       =   self.reshape_heads_to_batch_dim(lowerCamelCase__	)
  
  						if self.use_memory_efficient_attention:
  							UpperCAmelCase__:  str		       =   query_states.transpose(1							,			0							,			2	)
  							UpperCAmelCase__:  Any		       =   key_states.transpose(1							,			0							,			2	)
  							UpperCAmelCase__:  int		       =   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
  
  							UpperCAmelCase__:  List[Any]		       =   query_states.shape[-3]
  							if flatten_latent_dim % 6_4 == 0:
  								UpperCAmelCase__:  Optional[int]		       =   int(flatten_latent_dim / 6_4	)
  							elif flatten_latent_dim % 1_6 == 0:
  								UpperCAmelCase__:  List[str]		       =   int(flatten_latent_dim / 1_6	)
  							elif flatten_latent_dim % 4 == 0:
  								UpperCAmelCase__:  int		       =   int(flatten_latent_dim / 4	)
  							else:
  								UpperCAmelCase__:  str		       =   int(lowerCamelCase__	)
  
  							UpperCAmelCase__:  str		       =   jax_memory_efficient_attention(
  							    lowerCamelCase__							,			lowerCamelCase__							,			lowerCamelCase__							,			query_chunk_size=lowerCamelCase__							,			key_chunk_size=4_0_9_6 * 4	)
  
  							UpperCAmelCase__:  Union[str, Any]		       =   hidden_states.transpose(1							,			0							,			2	)
  						else:
  							# compute attentions
  							UpperCAmelCase__:  Union[str, Any]		       =   jnp.einsum("b i d, b j d->b i j"							,			lowerCamelCase__							,			lowerCamelCase__	)
  							UpperCAmelCase__:  str		       =   attention_scores * self.scale
  							UpperCAmelCase__:  Any		       =   nn.softmax(lowerCamelCase__							,			axis=2	)
  
  							# attend to values
  							UpperCAmelCase__:  Any		       =   jnp.einsum("b i j, b j d -> b i d"							,			lowerCamelCase__							,			lowerCamelCase__	)
  
  						UpperCAmelCase__:  int		       =   self.reshape_batch_dim_to_heads(lowerCamelCase__	)
  						UpperCAmelCase__:  List[str]		       =   self.proj_attn(lowerCamelCase__	)
  						return self.dropout_layer(lowerCamelCase__							,			deterministic=lowerCamelCase__	)
  
  class __UpperCamelCase    (							nn.Module		):
  
  
  
  
  
  
  					'''simple docstring'''
  
  
  
  					__magic_name__     			=      42
  					__magic_name__     			=      42
  					__magic_name__     			=      42
  					__magic_name__     			=      0.0
  					__magic_name__     			=      False
  					__magic_name__     			=      jnp.floataa
  					__magic_name__     			=      False
  
  
  
  
  
  					def       _UpperCAmelCase						(							self	):
  						# self attention (or cross_attention if only_cross_attention is True)
  						UpperCAmelCase__:  Any		       =   FlaxAttention(
  						    self.dim							,			self.n_heads							,			self.d_head							,			self.dropout							,			self.use_memory_efficient_attention							,			dtype=self.dtype	)
  						# cross attention
  						UpperCAmelCase__:  Dict		       =   FlaxAttention(
  						    self.dim							,			self.n_heads							,			self.d_head							,			self.dropout							,			self.use_memory_efficient_attention							,			dtype=self.dtype	)
  						UpperCAmelCase__:  List[str]		       =   FlaxFeedForward(dim=self.dim							,			dropout=self.dropout							,			dtype=self.dtype	)
  						UpperCAmelCase__:  Optional[Any]		       =   nn.LayerNorm(epsilon=1e-5							,			dtype=self.dtype	)
  						UpperCAmelCase__:  Optional[Any]		       =   nn.LayerNorm(epsilon=1e-5							,			dtype=self.dtype	)
  						UpperCAmelCase__:  List[str]		       =   nn.LayerNorm(epsilon=1e-5							,			dtype=self.dtype	)
  						UpperCAmelCase__:  Optional[int]		       =   nn.Dropout(rate=self.dropout	)
  
  
  
  
  
  					def __call__(							self							,			lowerCamelCase__							,			lowerCamelCase__							,			lowerCamelCase__=True	):
  						# self attention
  						UpperCAmelCase__:  Optional[Any]		       =   hidden_states
  						if self.only_cross_attention:
  							UpperCAmelCase__:  Optional[Any]		       =   self.attna(self.norma(lowerCamelCase__	)							,			lowerCamelCase__							,			deterministic=lowerCamelCase__	)
  						else:
  							UpperCAmelCase__:  List[str]		       =   self.attna(self.norma(lowerCamelCase__	)							,			deterministic=lowerCamelCase__	)
  						UpperCAmelCase__:  str		       =   hidden_states + residual
  
  						# cross attention
  						UpperCAmelCase__:  List[str]		       =   hidden_states
  						UpperCAmelCase__:  Dict		       =   self.attna(self.norma(lowerCamelCase__	)							,			lowerCamelCase__							,			deterministic=lowerCamelCase__	)
  						UpperCAmelCase__:  List[Any]		       =   hidden_states + residual
  
  						# feed forward
  						UpperCAmelCase__:  List[str]		       =   hidden_states
  						UpperCAmelCase__:  Optional[int]		       =   self.ff(self.norma(lowerCamelCase__	)							,			deterministic=lowerCamelCase__	)
  						UpperCAmelCase__:  Optional[Any]		       =   hidden_states + residual
  
  						return self.dropout_layer(lowerCamelCase__							,			deterministic=lowerCamelCase__	)
  
  class __UpperCamelCase    (							nn.Module		):
  
  
  
  
  
  
  					'''simple docstring'''
  
  
  
  					__magic_name__     			=      42
  					__magic_name__     			=      42
  					__magic_name__     			=      42
  					__magic_name__     			=      1
  					__magic_name__     			=      0.0
  					__magic_name__     			=      False
  					__magic_name__     			=      False
  					__magic_name__     			=      jnp.floataa
  					__magic_name__     			=      False
  
  
  
  
  
  					def       _UpperCAmelCase						(							self	):
  						UpperCAmelCase__:  Any		       =   nn.GroupNorm(num_groups=3_2							,			epsilon=1e-5	)
  
  						UpperCAmelCase__:  Optional[Any]		       =   self.n_heads * self.d_head
  						if self.use_linear_projection:
  							UpperCAmelCase__:  Union[str, Any]		       =   nn.Dense(lowerCamelCase__							,			dtype=self.dtype	)
  						else:
  							UpperCAmelCase__:  str		       =   nn.Conv(
  							    lowerCamelCase__							,			kernel_size=(1, 1)							,			strides=(1, 1)							,			padding="VALID"							,			dtype=self.dtype							,			)
  
  						UpperCAmelCase__:  int		       =   [
  						    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:
  							UpperCAmelCase__:  Tuple		       =   nn.Dense(lowerCamelCase__							,			dtype=self.dtype	)
  						else:
  							UpperCAmelCase__:  List[str]		       =   nn.Conv(
  							    lowerCamelCase__							,			kernel_size=(1, 1)							,			strides=(1, 1)							,			padding="VALID"							,			dtype=self.dtype							,			)
  
  						UpperCAmelCase__:  Optional[int]		       =   nn.Dropout(rate=self.dropout	)
  
  
  
  
  
  					def __call__(							self							,			lowerCamelCase__							,			lowerCamelCase__							,			lowerCamelCase__=True	):
  						UpperCAmelCase__       ,	UpperCAmelCase__       ,	UpperCAmelCase__       ,	UpperCAmelCase__:  Optional[int]		       =   hidden_states.shape
  						UpperCAmelCase__:  List[str]		       =   hidden_states
  						UpperCAmelCase__:  Optional[Any]		       =   self.norm(lowerCamelCase__	)
  						if self.use_linear_projection:
  							UpperCAmelCase__:  Union[str, Any]		       =   hidden_states.reshape(lowerCamelCase__							,			height * width							,			lowerCamelCase__	)
  							UpperCAmelCase__:  List[Any]		       =   self.proj_in(lowerCamelCase__	)
  						else:
  							UpperCAmelCase__:  Any		       =   self.proj_in(lowerCamelCase__	)
  							UpperCAmelCase__:  Tuple		       =   hidden_states.reshape(lowerCamelCase__							,			height * width							,			lowerCamelCase__	)
  
  						for transformer_block in self.transformer_blocks:
  							UpperCAmelCase__:  Any		       =   transformer_block(lowerCamelCase__							,			lowerCamelCase__							,			deterministic=lowerCamelCase__	)
  
  						if self.use_linear_projection:
  							UpperCAmelCase__:  Dict		       =   self.proj_out(lowerCamelCase__	)
  							UpperCAmelCase__:  Any		       =   hidden_states.reshape(lowerCamelCase__							,			lowerCamelCase__							,			lowerCamelCase__							,			lowerCamelCase__	)
  						else:
  							UpperCAmelCase__:  Union[str, Any]		       =   hidden_states.reshape(lowerCamelCase__							,			lowerCamelCase__							,			lowerCamelCase__							,			lowerCamelCase__	)
  							UpperCAmelCase__:  List[Any]		       =   self.proj_out(lowerCamelCase__	)
  
  						UpperCAmelCase__:  Dict		       =   hidden_states + residual
  						return self.dropout_layer(lowerCamelCase__							,			deterministic=lowerCamelCase__	)
  
  class __UpperCamelCase    (							nn.Module		):
  
  
  
  
  
  
  					'''simple docstring'''
  
  
  
  					__magic_name__     			=      42
  					__magic_name__     			=      0.0
  					__magic_name__     			=      jnp.floataa
  
  
  
  
  
  					def       _UpperCAmelCase						(							self	):
  						# The second linear layer needs to be called
  						# net_2 for now to match the index of the Sequential layer
  						UpperCAmelCase__:  Dict		       =   FlaxGEGLU(self.dim							,			self.dropout							,			self.dtype	)
  						UpperCAmelCase__:  int		       =   nn.Dense(self.dim							,			dtype=self.dtype	)
  
  
  
  
  
  					def __call__(							self							,			lowerCamelCase__							,			lowerCamelCase__=True	):
  						UpperCAmelCase__:  List[str]		       =   self.net_a(lowerCamelCase__							,			deterministic=lowerCamelCase__	)
  						UpperCAmelCase__:  List[str]		       =   self.net_a(lowerCamelCase__	)
  						return hidden_states
  
  class __UpperCamelCase    (							nn.Module		):
  
  
  
  
  
  
  					'''simple docstring'''
  
  
  
  					__magic_name__     			=      42
  					__magic_name__     			=      0.0
  					__magic_name__     			=      jnp.floataa
  
  
  
  
  
  					def       _UpperCAmelCase						(							self	):
  						UpperCAmelCase__:  List[Any]		       =   self.dim * 4
  						UpperCAmelCase__:  Any		       =   nn.Dense(inner_dim * 2							,			dtype=self.dtype	)
  						UpperCAmelCase__:  Dict		       =   nn.Dropout(rate=self.dropout	)
  
  
  
  
  
  					def __call__(							self							,			lowerCamelCase__							,			lowerCamelCase__=True	):
  						UpperCAmelCase__:  Optional[int]		       =   self.proj(lowerCamelCase__	)
  						UpperCAmelCase__       ,	UpperCAmelCase__:  List[Any]		       =   jnp.split(lowerCamelCase__							,			2							,			axis=2	)
  						return self.dropout_layer(hidden_linear * nn.gelu(lowerCamelCase__	)							,			deterministic=lowerCamelCase__	)
style_context: 
  
  import warnings
  
  import numpy as np
  import pandas as pd
  from matplotlib import pyplot as plt
  from sklearn.metrics import pairwise_distances
  
  warnings.filterwarnings("""ignore""")
  
  _lowerCAmelCase     :			str    ="""K-MEANS-CLUST/ """
  
  
  
  def  _A  (			SCREAMING_SNAKE_CASE  ,SCREAMING_SNAKE_CASE  ,SCREAMING_SNAKE_CASE=None			):
  	if seed is not None:  # useful for obtaining consistent results
  		np.random.seed(SCREAMING_SNAKE_CASE			)
  	UpperCAmelCase__:  Optional[Any]		       =   data.shape[0]  # number of data points
  
  	# Pick K indices from range [0, N).
  	UpperCAmelCase__:  Optional[int]		       =   np.random.randint(0  ,SCREAMING_SNAKE_CASE  ,SCREAMING_SNAKE_CASE			)
  
  	# Keep centroids as dense format, as many entries will be nonzero due to averaging.
  	# As long as at least one document in a cluster contains a word,
  	# it will carry a nonzero weight in the TF-IDF vector of the centroid.
  	UpperCAmelCase__:  Tuple		       =   data[rand_indices, :]
  
  	return centroids
  
  
  
  def  _A  (			SCREAMING_SNAKE_CASE  ,SCREAMING_SNAKE_CASE			):
  	return pairwise_distances(SCREAMING_SNAKE_CASE  ,SCREAMING_SNAKE_CASE  ,metric="euclidean"			)
  
  
  
  def  _A  (			SCREAMING_SNAKE_CASE  ,SCREAMING_SNAKE_CASE			):
  	# Compute distances between each data point and the set of centroids:
  	# Fill in the blank (RHS only)
  	UpperCAmelCase__:  Union[str, Any]		       =   centroid_pairwise_dist(SCREAMING_SNAKE_CASE  ,SCREAMING_SNAKE_CASE			)
  
  	# Compute cluster assignments for each data point:
  	# Fill in the blank (RHS only)
  	UpperCAmelCase__:  Dict		       =   np.argmin(SCREAMING_SNAKE_CASE  ,axis=1			)
  
  	return cluster_assignment
  
  
  
  def  _A  (			SCREAMING_SNAKE_CASE  ,SCREAMING_SNAKE_CASE  ,SCREAMING_SNAKE_CASE			):
  	UpperCAmelCase__:  str		       =   []
  	for i in range(SCREAMING_SNAKE_CASE			):
  		# Select all data points that belong to cluster i. Fill in the blank (RHS only)
  		UpperCAmelCase__:  Tuple		       =   data[cluster_assignment == i]
  		# Compute the mean of the data points. Fill in the blank (RHS only)
  		UpperCAmelCase__:  Any		       =   member_data_points.mean(axis=0			)
  		new_centroids.append(SCREAMING_SNAKE_CASE			)
  	UpperCAmelCase__:  List[Any]		       =   np.array(SCREAMING_SNAKE_CASE			)
  
  	return new_centroids
  
  
  
  def  _A  (			SCREAMING_SNAKE_CASE  ,SCREAMING_SNAKE_CASE  ,SCREAMING_SNAKE_CASE  ,SCREAMING_SNAKE_CASE			):
  	UpperCAmelCase__:  Tuple		       =   0.0
  	for i in range(SCREAMING_SNAKE_CASE			):
  		# Select all data points that belong to cluster i. Fill in the blank (RHS only)
  		UpperCAmelCase__:  Optional[Any]		       =   data[cluster_assignment == i, :]
  
  		if member_data_points.shape[0] > 0:  # check if i-th cluster is non-empty
  			# Compute distances from centroid to data points (RHS only)
  			UpperCAmelCase__:  int		       =   pairwise_distances(
  			    SCREAMING_SNAKE_CASE  ,[centroids[i]]  ,metric="euclidean"			)
  			UpperCAmelCase__:  int		       =   distances**2
  			heterogeneity += np.sum(SCREAMING_SNAKE_CASE			)
  
  	return heterogeneity
  
  
  
  def  _A  (			SCREAMING_SNAKE_CASE  ,SCREAMING_SNAKE_CASE			):
  	plt.figure(figsize=(7, 4)			)
  	plt.plot(SCREAMING_SNAKE_CASE  ,linewidth=4			)
  	plt.xlabel("# Iterations"			)
  	plt.ylabel("Heterogeneity"			)
  	plt.title(f"Heterogeneity of clustering over time, K={k:d}"			)
  	plt.rcParams.update({"font.size": 1_6}			)
  	plt.show()
  
  
  
  def  _A  (			SCREAMING_SNAKE_CASE  ,SCREAMING_SNAKE_CASE  ,SCREAMING_SNAKE_CASE  ,SCREAMING_SNAKE_CASE=5_0_0  ,SCREAMING_SNAKE_CASE=None  ,SCREAMING_SNAKE_CASE=False			):
  	UpperCAmelCase__:  Optional[int]		       =   initial_centroids[:]
  	UpperCAmelCase__:  List[str]		       =   None
  
  	for itr in range(SCREAMING_SNAKE_CASE			):
  		if verbose:
  			print(SCREAMING_SNAKE_CASE  ,end=""			)
  
  		# 1. Make cluster assignments using nearest centroids
  		UpperCAmelCase__:  Optional[Any]		       =   assign_clusters(SCREAMING_SNAKE_CASE  ,SCREAMING_SNAKE_CASE			)
  
  		# 2. Compute a new centroid for each of the k clusters, averaging all data
  		#    points assigned to that cluster.
  		UpperCAmelCase__:  Union[str, Any]		       =   revise_centroids(SCREAMING_SNAKE_CASE  ,SCREAMING_SNAKE_CASE  ,SCREAMING_SNAKE_CASE			)
  
  		# Check for convergence: if none of the assignments changed, stop
  		if (
  		    prev_cluster_assignment is not None
  		    and (prev_cluster_assignment == cluster_assignment).all()
  		):
  			break
  
  		# Print number of new assignments
  		if prev_cluster_assignment is not None:
  			UpperCAmelCase__:  Union[str, Any]		       =   np.sum(prev_cluster_assignment != cluster_assignment			)
  			if verbose:
  				print(
  				    f"    {num_changed:5d} elements changed their cluster assignment."			)
  
          # Record heterogeneity convergence metric
  		if record_heterogeneity is not None:
  			# YOUR CODE HERE
  			UpperCAmelCase__:  str		       =   compute_heterogeneity(SCREAMING_SNAKE_CASE  ,SCREAMING_SNAKE_CASE  ,SCREAMING_SNAKE_CASE  ,SCREAMING_SNAKE_CASE			)
  			record_heterogeneity.append(SCREAMING_SNAKE_CASE			)
  
  		UpperCAmelCase__:  str		       =   cluster_assignment[:]
  
  	return centroids, cluster_assignment
  
  
  # Mock test below
  if False:  # change to true to run this test case.
  					from sklearn import datasets as ds
  
  					_lowerCAmelCase     :			Optional[Any]    =ds.load_iris()
  					_lowerCAmelCase     :			int    =3
  					_lowerCAmelCase     :			List[str]    =[]
  					_lowerCAmelCase     :			int    =get_initial_centroids(dataset["""data"""], k, seed=0)
  					_lowerCAmelCase   ,  _lowerCAmelCase     :			Union[str, Any]    =kmeans(
  					    dataset["""data"""],
  					    k,
  					    initial_centroids,
  					    maxiter=4_00,
  					    record_heterogeneity=heterogeneity,
  					    verbose=True,
  					)
  					plot_heterogeneity(heterogeneity, k)
  
  
  
  def  _A  (			SCREAMING_SNAKE_CASE  ,SCREAMING_SNAKE_CASE  ,SCREAMING_SNAKE_CASE=None			):
  	# Fill missing values with given rules
  	if fill_missing_report:
  		UpperCAmelCase__:  int		       =   df.fillna(value=SCREAMING_SNAKE_CASE			)
  	UpperCAmelCase__:  Any		       =   1
  	UpperCAmelCase__:  Any		       =   df.select_dtypes(np.number			).columns
  	UpperCAmelCase__:  List[Any]		       =   (
  	    df.groupby(["Cluster"]			)[  # construct report dataframe
  	        numeric_cols
  	    ]  # group by cluster number
  	    .agg(
  	        [
  	            ("sum", np.sum),
  	            ("mean_with_zeros", lambda SCREAMING_SNAKE_CASE			: np.mean(np.nan_to_num(SCREAMING_SNAKE_CASE			)			)),
  	            ("mean_without_zeros", lambda SCREAMING_SNAKE_CASE			: x.replace(0  ,np.NaN			).mean()),
  	            (
  	                "mean_25-75",
  	                lambda SCREAMING_SNAKE_CASE			: np.mean(
  	                    np.nan_to_num(
  	                        sorted(SCREAMING_SNAKE_CASE			)[
  	                            round(len(SCREAMING_SNAKE_CASE			) * 2_5 / 1_0_0			) : round(len(SCREAMING_SNAKE_CASE			) * 7_5 / 1_0_0			)
  	                        ]			)			),
  	            ),
  	            ("mean_with_na", np.mean),
  	            ("min", lambda SCREAMING_SNAKE_CASE			: x.min()),
  	            ("5%", lambda SCREAMING_SNAKE_CASE			: x.quantile(0.05			)),
  	            ("25%", lambda SCREAMING_SNAKE_CASE			: x.quantile(0.25			)),
  	            ("50%", lambda SCREAMING_SNAKE_CASE			: x.quantile(0.50			)),
  	            ("75%", lambda SCREAMING_SNAKE_CASE			: x.quantile(0.75			)),
  	            ("95%", lambda SCREAMING_SNAKE_CASE			: x.quantile(0.95			)),
  	            ("max", lambda SCREAMING_SNAKE_CASE			: x.max()),
  	            ("count", lambda SCREAMING_SNAKE_CASE			: x.count()),
  	            ("stdev", lambda SCREAMING_SNAKE_CASE			: x.std()),
  	            ("mode", lambda SCREAMING_SNAKE_CASE			: x.mode()[0]),
  	            ("median", lambda SCREAMING_SNAKE_CASE			: x.median()),
  	            ("# > 0", lambda SCREAMING_SNAKE_CASE			: (x > 0).sum()),
  	        ]			)
  	    .T.reset_index()
  	    .rename(index=SCREAMING_SNAKE_CASE  ,columns={"level_0": "Features", "level_1": "Type"}			)
  	)  # rename columns
  	# calculate the size of cluster(count of clientID's)
  	UpperCAmelCase__:  str		       =   report[
  	    (report["Features"] == "dummy") & (report["Type"] == "count")
  	].copy()  # avoid SettingWithCopyWarning
  	UpperCAmelCase__:  int		       =   (
  	    "ClusterSize"  # rename created cluster df to match report column names
  	)
  	UpperCAmelCase__:  List[str]		       =   "# of Customers"
  	UpperCAmelCase__:  int		       =   pd.DataFrame(
  	    clustersize.iloc[:, 2:].values
  	    / clustersize.iloc[:, 2:].values.sum()			)
  	UpperCAmelCase__:  Dict		       =   "% of Customers"  # rename created cluster df to match report column names
  	UpperCAmelCase__:  Any		       =   "ClusterProportion"
  	UpperCAmelCase__:  int		       =   clusterproportion.columns.tolist()
  	UpperCAmelCase__:  int		       =   cols[-2:] + cols[:-2]
  	UpperCAmelCase__:  Optional[int]		       =   clusterproportion[cols]  # rearrange columns to match report
  	UpperCAmelCase__:  Union[str, Any]		       =   report.columns
  	UpperCAmelCase__:  int		       =   pd.DataFrame(
  	    abs(
  	        report[report["Type"] == "count"].iloc[:, 2:].values
  	        - clustersize.iloc[:, 2:].values			)			)  # generating df with count of nan values
  	UpperCAmelCase__:  Optional[int]		       =   0
  	UpperCAmelCase__:  Tuple		       =   "# of nan"
  	UpperCAmelCase__:  Any		       =   report[
  	    report["Type"] == "count"
  	].Features.tolist()  # filling values in order to match report
  	UpperCAmelCase__:  int		       =   a.columns.tolist()
  	UpperCAmelCase__:  Tuple		       =   cols[-2:] + cols[:-2]
  	UpperCAmelCase__:  List[str]		       =   a[cols]  # rearrange columns to match report
  	UpperCAmelCase__:  int		       =   report.columns  # rename columns to match report
  	UpperCAmelCase__:  List[Any]		       =   report.drop(
  	    report[report.Type == "count"].index			)  # drop count values except cluster size
  	UpperCAmelCase__:  Optional[Any]		       =   pd.concat(
  	    [report, a, clustersize, clusterproportion]  ,axis=0			)  # concat report with clustert size and nan values
  	UpperCAmelCase__:  Any		       =   report["Features"].isin(SCREAMING_SNAKE_CASE			)
  	UpperCAmelCase__:  Tuple		       =   report.columns.tolist()
  	UpperCAmelCase__:  str		       =   cols[0:2] + cols[-1:] + cols[2:-1]
  	UpperCAmelCase__:  int		       =   report[cols]
  	UpperCAmelCase__:  str		       =   {
  	    "ClusterSize": 9,
  	    "ClusterProportion": 8,
  	    "mean_with_zeros": 7,
  	    "mean_with_na": 6,
  	    "max": 5,
  	    "50%": 4,
  	    "min": 3,
  	    "25%": 2,
  	    "75%": 1,
  	    "# of nan": 0,
  	    "# > 0": -1,
  	    "sum_with_na": -2,
  	}
  	UpperCAmelCase__:  int		       =   (
  	    report.assign(
  	        Sortera=lambda SCREAMING_SNAKE_CASE			: x.Type.map(SCREAMING_SNAKE_CASE			)  ,Sortera=lambda SCREAMING_SNAKE_CASE			: list(reversed(range(len(SCREAMING_SNAKE_CASE			)			)			)			)  ,)
  	    .sort_values(["Sorter1", "Mark", "Sorter2"]  ,ascending=SCREAMING_SNAKE_CASE			)
  	    .drop(["Sorter1", "Sorter2"]  ,axis=1			)
  	)
  	UpperCAmelCase__:  List[str]		       =   ""
  	UpperCAmelCase__:  Any		       =   report.reset_index()
  	UpperCAmelCase__:  Any		       =   report.drop(columns=["index"]			)
  	return report
  
  
  if __name__ == "__main__":
  					import doctest
  
  					doctest.testmod()