code: 
  
  import unittest
  
  import numpy as np
  import timeout_decorator  # noqa
  
  from transformers import MarianConfig, is_flax_available
  from transformers.testing_utils import require_flax, require_sentencepiece, require_tokenizers, slow
  from transformers.utils import cached_property
  
  from ...generation.test_flax_utils import FlaxGenerationTesterMixin
  from ...test_modeling_flax_common import FlaxModelTesterMixin, ids_tensor
  
  
  if is_flax_available():
  						import os
  
  						# The slow tests are often failing with OOM error on GPU
  						# This makes JAX allocate exactly what is needed on demand, and deallocate memory that is no longer needed
  						# but will be slower as stated here https://jax.readthedocs.io/en/latest/gpu_memory_allocation.html
  						_lowerCAmelCase     :			List[Any]    ='platform'
  
  						import jax
  						import jax.numpy as jnp
  
  						from transformers import MarianTokenizer
  						from transformers.models.marian.modeling_flax_marian import FlaxMarianModel, FlaxMarianMTModel, shift_tokens_right
  
  
  
  def  _A  (			SCREAMING_SNAKE_CASE  ,SCREAMING_SNAKE_CASE  ,SCREAMING_SNAKE_CASE=None  ,SCREAMING_SNAKE_CASE=None  ,SCREAMING_SNAKE_CASE=None  ,SCREAMING_SNAKE_CASE=None  ,SCREAMING_SNAKE_CASE=None  ,SCREAMING_SNAKE_CASE=None  ,):
  
  
  
  	if attention_mask is None:
  		UpperCAmelCase__:  List[str]		       =   np.where(input_ids != config.pad_token_id  ,1  ,0			)
  	if decoder_attention_mask is None:
  		UpperCAmelCase__:  Dict		       =   np.where(decoder_input_ids != config.pad_token_id  ,1  ,0			)
  	if head_mask is None:
  		UpperCAmelCase__:  List[str]		       =   np.ones((config.encoder_layers, config.encoder_attention_heads)			)
  	if decoder_head_mask is None:
  		UpperCAmelCase__:  Union[str, Any]		       =   np.ones((config.decoder_layers, config.decoder_attention_heads)			)
  	if cross_attn_head_mask is None:
  		UpperCAmelCase__:  Optional[int]		       =   np.ones((config.decoder_layers, config.decoder_attention_heads)			)
  	return {
  	    "input_ids": input_ids,
  	    "decoder_input_ids": decoder_input_ids,
  	    "attention_mask": attention_mask,
  	    "decoder_attention_mask": attention_mask,
  	}
  
  class __UpperCamelCase    :
  
  
  
  
  
  
  					'''simple docstring'''
  
  
  
  					def __init__(							self							,			lowerCamelCase__							,			lowerCamelCase__=1_3							,			lowerCamelCase__=7							,			lowerCamelCase__=True							,			lowerCamelCase__=False							,			lowerCamelCase__=9_9							,			lowerCamelCase__=1_6							,			lowerCamelCase__=2							,			lowerCamelCase__=4							,			lowerCamelCase__=4							,			lowerCamelCase__="gelu"							,			lowerCamelCase__=0.1							,			lowerCamelCase__=0.1							,			lowerCamelCase__=3_2							,			lowerCamelCase__=2							,			lowerCamelCase__=1							,			lowerCamelCase__=0							,			lowerCamelCase__=0.02							,			):
  						UpperCAmelCase__:  List[str]		       =   parent
  						UpperCAmelCase__:  List[Any]		       =   batch_size
  						UpperCAmelCase__:  int		       =   seq_length
  						UpperCAmelCase__:  Dict		       =   is_training
  						UpperCAmelCase__:  Dict		       =   use_labels
  						UpperCAmelCase__:  str		       =   vocab_size
  						UpperCAmelCase__:  Optional[int]		       =   hidden_size
  						UpperCAmelCase__:  Optional[int]		       =   num_hidden_layers
  						UpperCAmelCase__:  Tuple		       =   num_attention_heads
  						UpperCAmelCase__:  List[str]		       =   intermediate_size
  						UpperCAmelCase__:  Optional[Any]		       =   hidden_act
  						UpperCAmelCase__:  str		       =   hidden_dropout_prob
  						UpperCAmelCase__:  Tuple		       =   attention_probs_dropout_prob
  						UpperCAmelCase__:  str		       =   max_position_embeddings
  						UpperCAmelCase__:  Optional[int]		       =   eos_token_id
  						UpperCAmelCase__:  List[str]		       =   pad_token_id
  						UpperCAmelCase__:  List[Any]		       =   bos_token_id
  						UpperCAmelCase__:  int		       =   initializer_range
  
  
  
  
  
  					def       _UpperCAmelCase						(							self	):
  						UpperCAmelCase__:  Optional[int]		       =   np.clip(ids_tensor([self.batch_size, self.seq_length - 1]							,			self.vocab_size	)							,			3							,			self.vocab_size	)
  						UpperCAmelCase__:  Optional[Any]		       =   np.concatenate((input_ids, 2 * np.ones((self.batch_size, 1)							,			dtype=np.intaa	))							,			-1	)
  
  						UpperCAmelCase__:  List[str]		       =   shift_tokens_right(lowerCamelCase__							,			1							,			2	)
  
  						UpperCAmelCase__:  str		       =   MarianConfig(
  						    vocab_size=self.vocab_size							,			d_model=self.hidden_size							,			encoder_layers=self.num_hidden_layers							,			decoder_layers=self.num_hidden_layers							,			encoder_attention_heads=self.num_attention_heads							,			decoder_attention_heads=self.num_attention_heads							,			encoder_ffn_dim=self.intermediate_size							,			decoder_ffn_dim=self.intermediate_size							,			dropout=self.hidden_dropout_prob							,			attention_dropout=self.attention_probs_dropout_prob							,			max_position_embeddings=self.max_position_embeddings							,			eos_token_id=self.eos_token_id							,			bos_token_id=self.bos_token_id							,			pad_token_id=self.pad_token_id							,			initializer_range=self.initializer_range							,			use_cache=lowerCamelCase__							,			)
  						UpperCAmelCase__:  str		       =   prepare_marian_inputs_dict(lowerCamelCase__							,			lowerCamelCase__							,			lowerCamelCase__	)
  						return config, inputs_dict
  
  
  
  
  
  					def       _UpperCAmelCase						(							self	):
  						UpperCAmelCase__       ,	UpperCAmelCase__:  Union[str, Any]		       =   self.prepare_config_and_inputs()
  						return config, inputs_dict
  
  
  
  
  
  					def       _UpperCAmelCase						(							self							,			lowerCamelCase__							,			lowerCamelCase__							,			lowerCamelCase__	):
  						UpperCAmelCase__:  int		       =   2_0
  						UpperCAmelCase__:  List[Any]		       =   model_class_name(lowerCamelCase__	)
  
  						UpperCAmelCase__:  Union[str, Any]		       =   model.encode(inputs_dict["input_ids"]	)
  
  						UpperCAmelCase__       ,	UpperCAmelCase__:  Dict		       =   (
  						    inputs_dict["decoder_input_ids"],
  						    inputs_dict["decoder_attention_mask"],
  						)
  
  						UpperCAmelCase__:  Any		       =   model.init_cache(decoder_input_ids.shape[0]							,			lowerCamelCase__							,			lowerCamelCase__	)
  						UpperCAmelCase__:  str		       =   jnp.ones((decoder_input_ids.shape[0], max_decoder_length)							,			dtype="i4"	)
  
  						UpperCAmelCase__:  Dict		       =   jnp.broadcast_to(
  						    jnp.arange(decoder_input_ids.shape[-1] - 1	)[None, :]							,			(decoder_input_ids.shape[0], decoder_input_ids.shape[-1] - 1)							,			)
  						UpperCAmelCase__:  Union[str, Any]		       =   model.decode(
  						    decoder_input_ids[:, :-1]							,			lowerCamelCase__							,			decoder_attention_mask=lowerCamelCase__							,			past_key_values=lowerCamelCase__							,			decoder_position_ids=lowerCamelCase__							,			)
  
  						UpperCAmelCase__:  Tuple		       =   jnp.array(decoder_input_ids.shape[0] * [[decoder_input_ids.shape[-1] - 1]]							,			dtype="i4"	)
  						UpperCAmelCase__:  int		       =   model.decode(
  						    decoder_input_ids[:, -1:]							,			lowerCamelCase__							,			decoder_attention_mask=lowerCamelCase__							,			past_key_values=outputs_cache.past_key_values							,			decoder_position_ids=lowerCamelCase__							,			)
  
  						UpperCAmelCase__:  str		       =   model.decode(lowerCamelCase__							,			lowerCamelCase__	)
  
  						UpperCAmelCase__:  Union[str, Any]		       =   np.max(np.abs((outputs_cache_next[0][:, -1, :5] - outputs[0][:, -1, :5])	)	)
  						self.parent.assertTrue(diff < 1e-3							,			msg=F"Max diff is {diff}"	)
  
  
  
  
  
  					def       _UpperCAmelCase						(							self							,			lowerCamelCase__							,			lowerCamelCase__							,			lowerCamelCase__	):
  						UpperCAmelCase__:  Dict		       =   2_0
  						UpperCAmelCase__:  Tuple		       =   model_class_name(lowerCamelCase__	)
  
  						UpperCAmelCase__:  List[Any]		       =   model.encode(inputs_dict["input_ids"]	)
  
  						UpperCAmelCase__       ,	UpperCAmelCase__:  List[str]		       =   (
  						    inputs_dict["decoder_input_ids"],
  						    inputs_dict["decoder_attention_mask"],
  						)
  
  						UpperCAmelCase__:  Optional[int]		       =   jnp.concatenate(
  						    [
  						        decoder_attention_mask,
  						        jnp.zeros((decoder_attention_mask.shape[0], max_decoder_length - decoder_attention_mask.shape[1])	),
  						    ]							,			axis=-1							,			)
  
  						UpperCAmelCase__:  Dict		       =   model.init_cache(decoder_input_ids.shape[0]							,			lowerCamelCase__							,			lowerCamelCase__	)
  						UpperCAmelCase__:  List[Any]		       =   jnp.broadcast_to(
  						    jnp.arange(decoder_input_ids.shape[-1] - 1	)[None, :]							,			(decoder_input_ids.shape[0], decoder_input_ids.shape[-1] - 1)							,			)
  
  						UpperCAmelCase__:  Tuple		       =   model.decode(
  						    decoder_input_ids[:, :-1]							,			lowerCamelCase__							,			decoder_attention_mask=lowerCamelCase__							,			past_key_values=lowerCamelCase__							,			decoder_position_ids=lowerCamelCase__							,			)
  						UpperCAmelCase__:  Any		       =   jnp.array(decoder_input_ids.shape[0] * [[decoder_input_ids.shape[-1] - 1]]							,			dtype="i4"	)
  						UpperCAmelCase__:  int		       =   model.decode(
  						    decoder_input_ids[:, -1:]							,			lowerCamelCase__							,			past_key_values=outputs_cache.past_key_values							,			decoder_attention_mask=lowerCamelCase__							,			decoder_position_ids=lowerCamelCase__							,			)
  
  						UpperCAmelCase__:  str		       =   model.decode(lowerCamelCase__							,			lowerCamelCase__							,			decoder_attention_mask=lowerCamelCase__	)
  
  						UpperCAmelCase__:  Optional[int]		       =   np.max(np.abs((outputs_cache_next[0][:, -1, :5] - outputs[0][:, -1, :5])	)	)
  						self.parent.assertTrue(diff < 1e-3							,			msg=F"Max diff is {diff}"	)
  
  @require_flax
  class __UpperCamelCase    (							UpperCAmelCase_				,unittest.TestCase				,UpperCAmelCase_		):
  
  
  
  
  
  
  					'''simple docstring'''
  
  
  
  					__magic_name__     			=      True
  					__magic_name__     			=      (FlaxMarianModel, FlaxMarianMTModel) if is_flax_available() else ()
  					__magic_name__     			=      (FlaxMarianMTModel,) if is_flax_available() else ()
  
  
  
  
  
  					def       _UpperCAmelCase						(							self	):
  						UpperCAmelCase__:  Dict		       =   FlaxMarianModelTester(self	)
  
  
  
  
  
  					def       _UpperCAmelCase						(							self	):
  						UpperCAmelCase__       ,	UpperCAmelCase__:  Union[str, Any]		       =   self.model_tester.prepare_config_and_inputs()
  						for model_class in self.all_model_classes:
  							self.model_tester.check_use_cache_forward(lowerCamelCase__							,			lowerCamelCase__							,			lowerCamelCase__	)
  
  
  
  
  
  					def       _UpperCAmelCase						(							self	):
  						UpperCAmelCase__       ,	UpperCAmelCase__:  Optional[Any]		       =   self.model_tester.prepare_config_and_inputs()
  						for model_class in self.all_model_classes:
  							self.model_tester.check_use_cache_forward_with_attn_mask(lowerCamelCase__							,			lowerCamelCase__							,			lowerCamelCase__	)
  
  
  
  
  
  					def       _UpperCAmelCase						(							self	):
  						UpperCAmelCase__       ,	UpperCAmelCase__:  Dict		       =   self.model_tester.prepare_config_and_inputs_for_common()
  
  						for model_class in self.all_model_classes:
  							with self.subTest(model_class.__name__	):
  								UpperCAmelCase__:  Optional[Any]		       =   self._prepare_for_class(lowerCamelCase__							,			lowerCamelCase__	)
  								UpperCAmelCase__:  str		       =   model_class(lowerCamelCase__	)
  
  								@jax.jit
  								def encode_jitted(lowerCamelCase__							,			lowerCamelCase__=None							,			**lowerCamelCase__	):
  									return model.encode(input_ids=lowerCamelCase__							,			attention_mask=lowerCamelCase__	)
  
  								with self.subTest("JIT Enabled"	):
  									UpperCAmelCase__:  Dict		       =   encode_jitted(**lowerCamelCase__	).to_tuple()
  
  								with self.subTest("JIT Disabled"	):
  									with jax.disable_jit():
  										UpperCAmelCase__:  int		       =   encode_jitted(**lowerCamelCase__	).to_tuple()
  
  								self.assertEqual(len(lowerCamelCase__	)							,			len(lowerCamelCase__	)	)
  								for jitted_output, output in zip(lowerCamelCase__							,			lowerCamelCase__	):
  									self.assertEqual(jitted_output.shape							,			output.shape	)
  
  
  
  
  
  					def       _UpperCAmelCase						(							self	):
  						UpperCAmelCase__       ,	UpperCAmelCase__:  Union[str, Any]		       =   self.model_tester.prepare_config_and_inputs_for_common()
  
  						for model_class in self.all_model_classes:
  							with self.subTest(model_class.__name__	):
  								UpperCAmelCase__:  Tuple		       =   model_class(lowerCamelCase__	)
  								UpperCAmelCase__:  Any		       =   model.encode(inputs_dict["input_ids"]							,			inputs_dict["attention_mask"]	)
  
  								UpperCAmelCase__:  Union[str, Any]		       =   {
  								    "decoder_input_ids": inputs_dict["decoder_input_ids"],
  								    "decoder_attention_mask": inputs_dict["decoder_attention_mask"],
  								    "encoder_outputs": encoder_outputs,
  								}
  
  								@jax.jit
  								def decode_jitted(lowerCamelCase__							,			lowerCamelCase__							,			lowerCamelCase__	):
  									return model.decode(
  									    decoder_input_ids=lowerCamelCase__							,			decoder_attention_mask=lowerCamelCase__							,			encoder_outputs=lowerCamelCase__							,			)
  
  								with self.subTest("JIT Enabled"	):
  									UpperCAmelCase__:  Optional[int]		       =   decode_jitted(**lowerCamelCase__	).to_tuple()
  
  								with self.subTest("JIT Disabled"	):
  									with jax.disable_jit():
  										UpperCAmelCase__:  int		       =   decode_jitted(**lowerCamelCase__	).to_tuple()
  
  								self.assertEqual(len(lowerCamelCase__	)							,			len(lowerCamelCase__	)	)
  								for jitted_output, output in zip(lowerCamelCase__							,			lowerCamelCase__	):
  									self.assertEqual(jitted_output.shape							,			output.shape	)
  
  
  
  
  
  					@slow
  					def       _UpperCAmelCase						(							self	):
  						for model_class_name in self.all_model_classes:
  							UpperCAmelCase__:  Dict		       =   model_class_name.from_pretrained("Helsinki-NLP/opus-mt-en-de"	)
  							# FlaxMarianForSequenceClassification expects eos token in input_ids
  							UpperCAmelCase__:  int		       =   np.ones((1, 1)	) * model.config.eos_token_id
  							UpperCAmelCase__:  List[str]		       =   model(lowerCamelCase__	)
  							self.assertIsNotNone(lowerCamelCase__	)
  
  @require_flax
  @require_sentencepiece
  @require_tokenizers
  class __UpperCamelCase    (							unittest.TestCase		):
  
  
  
  
  
  
  					'''simple docstring'''
  
  
  
  					__magic_name__     			=      None
  					__magic_name__     			=      None
  
  
  
  
  
  					@classmethod
  					def       _UpperCAmelCase						(							cls	):
  						UpperCAmelCase__:  List[Any]		       =   F"Helsinki-NLP/opus-mt-{cls.src}-{cls.tgt}"
  						return cls
  
  
  
  
  
  					@cached_property
  					def       _UpperCAmelCase						(							self	):
  						return MarianTokenizer.from_pretrained(self.model_name	)
  
  
  
  
  
  					@property
  					def       _UpperCAmelCase						(							self	):
  						return self.tokenizer.eos_token_id
  
  
  
  
  
  					@cached_property
  					def       _UpperCAmelCase						(							self	):
  						UpperCAmelCase__:  int		       =   FlaxMarianMTModel.from_pretrained(self.model_name	)
  						self.assertEqual(model.config.decoder_start_token_id							,			model.config.pad_token_id	)
  						return model
  
  
  
  
  
  					def       _UpperCAmelCase						(							self							,			**lowerCamelCase__	):
  						UpperCAmelCase__:  Union[str, Any]		       =   self.translate_src_text(**lowerCamelCase__	)
  						self.assertListEqual(self.expected_text							,			lowerCamelCase__	)
  
  
  
  
  
  					def       _UpperCAmelCase						(							self							,			**lowerCamelCase__	):
  						UpperCAmelCase__:  int		       =   self.tokenizer(self.src_text							,			padding=lowerCamelCase__							,			return_tensors="np"							,			**lowerCamelCase__	)
  						UpperCAmelCase__:  List[str]		       =   self.model.generate(
  						    model_inputs.input_ids							,			attention_mask=model_inputs.attention_mask							,			num_beams=2							,			max_length=1_2_8							,			).sequences
  						UpperCAmelCase__:  List[Any]		       =   self.tokenizer.batch_decode(lowerCamelCase__							,			skip_special_tokens=lowerCamelCase__	)
  						return generated_words
  
  @require_flax
  @require_sentencepiece
  @require_tokenizers
  class __UpperCamelCase    (							UpperCAmelCase_		):
  
  
  
  
  
  
  					'''simple docstring'''
  
  
  
  					__magic_name__     			=      "en"
  					__magic_name__     			=      "fr"
  					__magic_name__     			=      [
  					    "I am a small frog.",
  					    "Now I can forget the 100 words of german that I know.",
  					]
  					__magic_name__     			=      [
  					    "Je suis une petite grenouille.",
  					    "Maintenant, je peux oublier les 100 mots d\'allemand que je connais.",
  					]
  
  
  
  
  
  					@slow
  					def       _UpperCAmelCase						(							self	):
  						self._assert_generated_batch_equal_expected()
  
  @require_flax
  @require_sentencepiece
  @require_tokenizers
  class __UpperCamelCase    (							UpperCAmelCase_		):
  
  
  
  
  
  
  					'''simple docstring'''
  
  
  
  					__magic_name__     			=      "fr"
  					__magic_name__     			=      "en"
  					__magic_name__     			=      [
  					    "Donnez moi le micro.",
  					    "Tom et Mary étaient assis à une table.",  # Accents
  					]
  					__magic_name__     			=      [
  					    "Give me the microphone.",
  					    "Tom and Mary were sitting at a table.",
  					]
  
  
  
  
  
  					@slow
  					def       _UpperCAmelCase						(							self	):
  						self._assert_generated_batch_equal_expected()
  
  @require_flax
  @require_sentencepiece
  @require_tokenizers
  class __UpperCamelCase    (							UpperCAmelCase_		):
  
  
  
  
  
  
  					'''simple docstring'''
  
  
  
  					__magic_name__     			=      "mt"
  					__magic_name__     			=      "en"
  					__magic_name__     			=      ["Billi messu b\'mod ġentili, Ġesù fejjaq raġel li kien milqut bil - marda kerha tal - ġdiem."]
  					__magic_name__     			=      ["Touching gently, Jesus healed a man who was affected by the sad disease of leprosy."]
  
  
  
  
  
  					@slow
  					def       _UpperCAmelCase						(							self	):
  						self._assert_generated_batch_equal_expected()
  
  @require_flax
  @require_sentencepiece
  @require_tokenizers
  class __UpperCamelCase    (							UpperCAmelCase_		):
  
  
  
  
  
  
  					'''simple docstring'''
  
  
  
  					__magic_name__     			=      "en"
  					__magic_name__     			=      "de"
  					__magic_name__     			=      [
  					    "I am a small frog.",
  					    "Now I can forget the 100 words of german that I know.",
  					    "Tom asked his teacher for advice.",
  					    "That\'s how I would do it.",
  					    "Tom really admired Mary\'s courage.",
  					    "Turn around and close your eyes.",
  					]
  					__magic_name__     			=      [
  					    "Ich bin ein kleiner Frosch.",
  					    "Jetzt kann ich die 100 Wörter des Deutschen vergessen, die ich kenne.",
  					    "Tom bat seinen Lehrer um Rat.",
  					    "So würde ich das machen.",
  					    "Tom bewunderte Marias Mut wirklich.",
  					    "Drehen Sie sich um und schließen Sie die Augen.",
  					]
  
  
  
  
  
  					@slow
  					def       _UpperCAmelCase						(							self	):
  						self._assert_generated_batch_equal_expected()
  
  @require_flax
  @require_sentencepiece
  @require_tokenizers
  class __UpperCamelCase    (							UpperCAmelCase_		):
  
  
  
  
  
  
  					'''simple docstring'''
  
  
  
  					__magic_name__     			=      "en"
  					__magic_name__     			=      "zh"
  					__magic_name__     			=      ["My name is Wolfgang and I live in Berlin"]
  					__magic_name__     			=      ["我叫沃尔夫冈 我住在柏林"]
  
  
  
  
  
  					@slow
  					def       _UpperCAmelCase						(							self	):
  						self._assert_generated_batch_equal_expected()
  
  @require_flax
  @require_sentencepiece
  @require_tokenizers
  class __UpperCamelCase    (							UpperCAmelCase_		):
  
  
  
  
  
  
  					'''simple docstring'''
  
  
  
  					__magic_name__     			=      "ru"
  					__magic_name__     			=      "fr"
  					__magic_name__     			=      ["Он показал мне рукопись своей новой пьесы."]
  					__magic_name__     			=      ["Il m\'a montré le manuscrit de sa nouvelle pièce."]
  
  
  
  
  
  					@slow
  					def       _UpperCAmelCase						(							self	):
  						self._assert_generated_batch_equal_expected()
  
  @require_flax
  @require_sentencepiece
  @require_tokenizers
  class __UpperCamelCase    (							UpperCAmelCase_		):
  
  
  
  
  
  
  					'''simple docstring'''
  
  
  
  					__magic_name__     			=      "en"
  					__magic_name__     			=      "ROMANCE"
  					__magic_name__     			=      [
  					    ">>fr<< Don\'t spend so much time watching TV.",
  					    ">>pt<< Your message has been sent.",
  					    ">>es<< He\'s two years older than me.",
  					]
  					__magic_name__     			=      [
  					    "Ne passez pas autant de temps à regarder la télé.",
  					    "A sua mensagem foi enviada.",
  					    "Es dos años más viejo que yo.",
  					]
  
  
  
  
  
  					@slow
  					def       _UpperCAmelCase						(							self	):
  						self._assert_generated_batch_equal_expected()
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()