




from random import shuffle

import tensorflow as tf
from numpy import array




def 		lowerCAmelCase_ (						lowerCamelCase						,  lowerCamelCase ):
				__magic_name__						:						List[str]										=int(lowerCamelCase )
				assert noofclusters < len(lowerCamelCase )

				# Find out the dimensionality
				__magic_name__						:						Union[str, Any]										=len(vectors[0] )

				# Will help select random centroids from among the available vectors
				__magic_name__						:						List[Any]										=list(range(len(lowerCamelCase ) ) )
				shuffle(lowerCamelCase )

				# GRAPH OF COMPUTATION
				# We initialize a new graph and set it as the default during each run
				# of this algorithm. This ensures that as this function is called
				# multiple times, the default graph doesn't keep getting crowded with
				# unused ops and Variables from previous function calls.

				__magic_name__						:						Any										=tf.Graph()

				with graph.as_default():
								# SESSION OF COMPUTATION

								__magic_name__						:						Any										=tf.Session()

								##CONSTRUCTING THE ELEMENTS OF COMPUTATION

								##First lets ensure we have a Variable vector for each centroid,
								##initialized to one of the vectors from the available data points
								__magic_name__						:						Dict										=[
								    tf.Variable(vectors[vector_indices[i]] ) for i in range(lowerCamelCase )
								]
								##These nodes will assign the centroid Variables the appropriate
								##values
								__magic_name__						:						List[Any]										=tf.placeholder("""float64"""						,  [dim] )
								__magic_name__						:						Optional[Any]										=[]
								for centroid in centroids:
												cent_assigns.append(tf.assign(lowerCamelCase						,  lowerCamelCase ) )

								##Variables for cluster assignments of individual vectors(initialized
								##to 0 at first)
								__magic_name__						:						Union[str, Any]										=[tf.Variable(0 ) for i in range(len(lowerCamelCase ) )]
								##These nodes will assign an assignment Variable the appropriate
								##value
								__magic_name__						:						Optional[Any]										=tf.placeholder("""int32""" )
								__magic_name__						:						List[str]										=[]
								for assignment in assignments:
												cluster_assigns.append(tf.assign(lowerCamelCase						,  lowerCamelCase ) )

								##Now lets construct the node that will compute the mean
								# The placeholder for the input
								__magic_name__						:						List[str]										=tf.placeholder("""float"""						,  [None, dim] )
								# The Node/op takes the input and computes a mean along the 0th
								# dimension, i.e. the list of input vectors
								__magic_name__						:						List[str]										=tf.reduce_mean(lowerCamelCase						,  0 )

								##Node for computing Euclidean distances
								# Placeholders for input
								__magic_name__						:						Union[str, Any]										=tf.placeholder("""float"""						,  [dim] )
								__magic_name__						:						Any										=tf.placeholder("""float"""						,  [dim] )
								__magic_name__						:						Any										=tf.sqrt(tf.reduce_sum(tf.pow(tf.sub(lowerCamelCase						,  lowerCamelCase )						,  2 ) ) )

								##This node will figure out which cluster to assign a vector to,
								##based on Euclidean distances of the vector from the centroids.
								# Placeholder for input
								__magic_name__						:						Any										=tf.placeholder("""float"""						,  [noofclusters] )
								__magic_name__						:						Tuple										=tf.argmin(lowerCamelCase						,  0 )

								##INITIALIZING STATE VARIABLES

								##This will help initialization of all Variables defined with respect
								##to the graph. The Variable-initializer should be defined after
								##all the Variables have been constructed, so that each of them
								##will be included in the initialization.
								__magic_name__						:						str										=tf.initialize_all_variables()

								# Initialize all variables
								sess.run(lowerCamelCase )

								##CLUSTERING ITERATIONS

								# Now perform the Expectation-Maximization steps of K-Means clustering
								# iterations. To keep things simple, we will only do a set number of
								# iterations, instead of using a Stopping Criterion.
								__magic_name__						:						List[str]										=100
								for _ in range(lowerCamelCase ):
												##EXPECTATION STEP
												##Based on the centroid locations till last iteration, compute
												##the _expected_ centroid assignments.
												# Iterate over each vector
												for vector_n in range(len(lowerCamelCase ) ):
																__magic_name__						:						List[str]										=vectors[vector_n]
																# Compute Euclidean distance between this vector and each
																# centroid. Remember that this list cannot be named
																#'centroid_distances', since that is the input to the
																# cluster assignment node.
																__magic_name__						:						Optional[Any]										=[
																    sess.run(lowerCamelCase						,  feed_dict={va: vect, va: sess.run(lowerCamelCase )} )
																    for centroid in centroids
																]
																# Now use the cluster assignment node, with the distances
																# as the input
																__magic_name__						:						Tuple										=sess.run(
																    lowerCamelCase						,  feed_dict={centroid_distances: distances} )
																# Now assign the value to the appropriate state variable
																sess.run(
																    cluster_assigns[vector_n]						,  feed_dict={assignment_value: assignment} )

												##MAXIMIZATION STEP
												# Based on the expected state computed from the Expectation Step,
												# compute the locations of the centroids so as to maximize the
												# overall objective of minimizing within-cluster Sum-of-Squares
												for cluster_n in range(lowerCamelCase ):
																# Collect all the vectors assigned to this cluster
																__magic_name__						:						List[str]										=[
																    vectors[i]
																    for i in range(len(lowerCamelCase ) )
																    if sess.run(assignments[i] ) == cluster_n
																]
																# Compute new centroid location
																__magic_name__						:						Tuple										=sess.run(
																    lowerCamelCase						,  feed_dict={mean_input: array(lowerCamelCase )} )
																# Assign value to appropriate variable
																sess.run(
																    cent_assigns[cluster_n]						,  feed_dict={centroid_value: new_location} )

        # Return centroids and assignments
								__magic_name__						:						List[Any]										=sess.run(lowerCamelCase )
								__magic_name__						:						int										=sess.run(lowerCamelCase )
								return centroids, assignments










import inspect
import unittest

import torch
import torch.nn as nn

from accelerate.hooks import (
    AlignDevicesHook,
    ModelHook,
    SequentialHook,
    add_hook_to_module,
    attach_align_device_hook,
    remove_hook_from_module,
    remove_hook_from_submodules,
)
from accelerate.test_utils import require_multi_gpu




class 		__A   (      nn.Module      ):

						def __init__(		self :List[Any]			):
										'''simple docstring'''

										super().__init__()
										__magic_name__						:						Tuple										=nn.Linear(3     ,			4			)
										__magic_name__						:						Union[str, Any]										=nn.BatchNormad(4			)
										__magic_name__						:						List[str]										=nn.Linear(4     ,			5			)


						def      A__    (		self :Dict     ,			__snake_case :Tuple			):
										'''simple docstring'''

										return self.lineara(self.batchnorm(self.lineara(__snake_case			)			)			)




class 		__A   (      UpperCamelCase__      ):

						def      A__    (		self :Any     ,			__snake_case :Optional[Any]     ,			*__snake_case :List[Any]     ,			**__snake_case :Any			):
										'''simple docstring'''

										return (args[0] + 1,) + args[1:], kwargs




class 		__A   (      UpperCamelCase__      ):

						def      A__    (		self :List[str]     ,			__snake_case :Tuple     ,			__snake_case :Union[str, Any]			):
										'''simple docstring'''

										return output + 1




class 		__A   (      unittest.TestCase      ):

						def      A__    (		self :List[str]			):
										'''simple docstring'''

										__magic_name__						:						int										=ModelForTest()
										__magic_name__						:						Tuple										=ModelHook()

										add_hook_to_module(__snake_case     ,			__snake_case			)
										self.assertEqual(test_model._hf_hook     ,			__snake_case			)
										self.assertTrue(hasattr(__snake_case     ,			"""_old_forward"""			)			)

										# Check adding the hook did not change the name or the signature
										self.assertEqual(test_model.forward.__name__     ,			"""forward"""			)
										self.assertListEqual(list(inspect.signature(test_model.forward			).parameters			)     ,			["""x"""]			)

										remove_hook_from_module(__snake_case			)
										self.assertFalse(hasattr(__snake_case     ,			"""_hf_hook"""			)			)
										self.assertFalse(hasattr(__snake_case     ,			"""_old_forward"""			)			)

						def      A__    (		self :Tuple			):
										'''simple docstring'''

										__magic_name__						:						int										=ModelForTest()
										__magic_name__						:						List[str]										=ModelHook()

										add_hook_to_module(__snake_case     ,			__snake_case			)
										add_hook_to_module(__snake_case     ,			__snake_case     ,			append=__snake_case			)

										self.assertEqual(isinstance(test_model._hf_hook     ,			__snake_case			)     ,			__snake_case			)
										self.assertEqual(len(test_model._hf_hook.hooks			)     ,			2			)
										self.assertTrue(hasattr(__snake_case     ,			"""_old_forward"""			)			)

										# Check adding the hook did not change the name or the signature
										self.assertEqual(test_model.forward.__name__     ,			"""forward"""			)
										self.assertListEqual(list(inspect.signature(test_model.forward			).parameters			)     ,			["""x"""]			)

										remove_hook_from_module(__snake_case			)
										self.assertFalse(hasattr(__snake_case     ,			"""_hf_hook"""			)			)
										self.assertFalse(hasattr(__snake_case     ,			"""_old_forward"""			)			)

						def      A__    (		self :List[str]			):
										'''simple docstring'''

										__magic_name__						:						Any										=ModelForTest()
										__magic_name__						:						Any										=torch.randn(2     ,			3			)
										__magic_name__						:						Any										=test_model(x + 1			)
										__magic_name__						:						Optional[Any]										=test_model(x + 2			)

										__magic_name__						:						int										=PreForwardHook()
										add_hook_to_module(__snake_case     ,			__snake_case			)
										__magic_name__						:						int										=test_model(__snake_case			)
										self.assertTrue(torch.allclose(__snake_case     ,			__snake_case     ,			atol=1E-5			)			)

										# Attaching a hook to a model when it already has one replaces, does not chain
										__magic_name__						:						str										=PreForwardHook()
										add_hook_to_module(__snake_case     ,			__snake_case			)
										__magic_name__						:						List[str]										=test_model(__snake_case			)
										self.assertTrue(torch.allclose(__snake_case     ,			__snake_case     ,			atol=1E-5			)			)

										# You need to use the sequential hook to chain two or more hooks
										__magic_name__						:						Optional[Any]										=SequentialHook(PreForwardHook()     ,			PreForwardHook()			)
										add_hook_to_module(__snake_case     ,			__snake_case			)

										__magic_name__						:						Any										=test_model(__snake_case			)
										assert torch.allclose(__snake_case     ,			__snake_case     ,			atol=1E-5			)

						def      A__    (		self :Any			):
										'''simple docstring'''

										__magic_name__						:						Optional[Any]										=ModelForTest()
										__magic_name__						:						Dict										=torch.randn(2     ,			3			)
										__magic_name__						:						Any										=test_model(__snake_case			)

										__magic_name__						:						Dict										=PostForwardHook()
										add_hook_to_module(__snake_case     ,			__snake_case			)
										__magic_name__						:						Any										=test_model(__snake_case			)
										self.assertTrue(torch.allclose(__snake_case     ,			output + 1     ,			atol=1E-5			)			)

										# Attaching a hook to a model when it already has one replaces, does not chain
										__magic_name__						:						Union[str, Any]										=PostForwardHook()
										add_hook_to_module(__snake_case     ,			__snake_case			)
										__magic_name__						:						Optional[int]										=test_model(__snake_case			)
										self.assertTrue(torch.allclose(__snake_case     ,			output + 1     ,			atol=1E-5			)			)

										# You need to use the sequential hook to chain two or more hooks
										__magic_name__						:						Union[str, Any]										=SequentialHook(PostForwardHook()     ,			PostForwardHook()			)
										add_hook_to_module(__snake_case     ,			__snake_case			)

										__magic_name__						:						Union[str, Any]										=test_model(__snake_case			)
										assert torch.allclose(__snake_case     ,			output + 2     ,			atol=1E-5			)

						def      A__    (		self :Tuple			):
										'''simple docstring'''

										__magic_name__						:						Tuple										=ModelForTest()
										__magic_name__						:						int										=torch.randn(2     ,			3			)
										__magic_name__						:						Union[str, Any]										=test_model(__snake_case			)

										__magic_name__						:						Union[str, Any]										=PostForwardHook()
										add_hook_to_module(__snake_case     ,			__snake_case			)
										__magic_name__						:						Dict										=test_model(__snake_case			)
										self.assertTrue(torch.allclose(__snake_case     ,			output + 1			)			)
										self.assertTrue(outputa.requires_grad			)

										__magic_name__						:						Any										=True
										__magic_name__						:						Any										=test_model(__snake_case			)
										self.assertFalse(outputa.requires_grad			)

						@require_multi_gpu
						def      A__    (		self :List[str]			):
										'''simple docstring'''

										__magic_name__						:						Optional[Any]										=ModelForTest()
										# Everything is on CPU
										self.assertEqual(model.lineara.weight.device     ,			torch.device("""cpu"""			)			)
										self.assertEqual(model.batchnorm.weight.device     ,			torch.device("""cpu"""			)			)
										self.assertEqual(model.lineara.weight.device     ,			torch.device("""cpu"""			)			)

										# This will move each submodule on different devices
										add_hook_to_module(model.lineara     ,			AlignDevicesHook(execution_device=0			)			)
										add_hook_to_module(model.batchnorm     ,			AlignDevicesHook(execution_device=0			)			)
										add_hook_to_module(model.lineara     ,			AlignDevicesHook(execution_device=1			)			)

										self.assertEqual(model.lineara.weight.device     ,			torch.device(0			)			)
										self.assertEqual(model.batchnorm.weight.device     ,			torch.device(0			)			)
										self.assertEqual(model.batchnorm.running_mean.device     ,			torch.device(0			)			)
										self.assertEqual(model.lineara.weight.device     ,			torch.device(1			)			)

										# We can still make a forward pass. The input does not need to be on any particular device
										__magic_name__						:						List[Any]										=torch.randn(2     ,			3			)
										__magic_name__						:						Optional[Any]										=model(__snake_case			)
										self.assertEqual(output.device     ,			torch.device(1			)			)

										# We can add a general hook to put back output on same device as input.
										add_hook_to_module(__snake_case     ,			AlignDevicesHook(io_same_device=__snake_case			)			)
										__magic_name__						:						int										=torch.randn(2     ,			3			).to(0			)
										__magic_name__						:						Optional[int]										=model(__snake_case			)
										self.assertEqual(output.device     ,			torch.device(0			)			)

						def      A__    (		self :List[Any]			):
										'''simple docstring'''

										__magic_name__						:						int										=ModelForTest()

										# Everything is on CPU
										self.assertEqual(model.lineara.weight.device     ,			torch.device("""cpu"""			)			)
										self.assertEqual(model.batchnorm.weight.device     ,			torch.device("""cpu"""			)			)
										self.assertEqual(model.lineara.weight.device     ,			torch.device("""cpu"""			)			)

										# This will move each submodule on different devices
										__magic_name__						:						int										={"""execution_device""": 0 if torch.cuda.is_available() else """cpu""", """offload""": True}

										add_hook_to_module(model.lineara     ,			AlignDevicesHook(**__snake_case			)			)
										add_hook_to_module(model.batchnorm     ,			AlignDevicesHook(**__snake_case			)			)
										add_hook_to_module(model.lineara     ,			AlignDevicesHook(**__snake_case			)			)

										# Parameters have been offloaded, so on the meta device
										self.assertEqual(model.lineara.weight.device     ,			torch.device("""meta"""			)			)
										self.assertEqual(model.batchnorm.weight.device     ,			torch.device("""meta"""			)			)
										self.assertEqual(model.lineara.weight.device     ,			torch.device("""meta"""			)			)
										# Buffers are not included in the offload by default, so are on the execution device
										__magic_name__						:						Optional[int]										=torch.device(hook_kwargs["""execution_device"""]			)
										self.assertEqual(model.batchnorm.running_mean.device     ,			__snake_case			)

										__magic_name__						:						Union[str, Any]										=torch.randn(2     ,			3			)
										__magic_name__						:						Optional[int]										=model(__snake_case			)
										self.assertEqual(output.device     ,			__snake_case			)

										# Removing hooks loads back the weights in the model.
										remove_hook_from_module(model.lineara			)
										remove_hook_from_module(model.batchnorm			)
										remove_hook_from_module(model.lineara			)
										self.assertEqual(model.lineara.weight.device     ,			torch.device("""cpu"""			)			)
										self.assertEqual(model.batchnorm.weight.device     ,			torch.device("""cpu"""			)			)
										self.assertEqual(model.lineara.weight.device     ,			torch.device("""cpu"""			)			)

										# Now test with buffers included in the offload
										__magic_name__						:						Tuple										={
										    """execution_device""": 0 if torch.cuda.is_available() else """cpu""",
										    """offload""": True,
										    """offload_buffers""": True,
										}

										add_hook_to_module(model.lineara     ,			AlignDevicesHook(**__snake_case			)			)
										add_hook_to_module(model.batchnorm     ,			AlignDevicesHook(**__snake_case			)			)
										add_hook_to_module(model.lineara     ,			AlignDevicesHook(**__snake_case			)			)

										# Parameters have been offloaded, so on the meta device, buffers included
										self.assertEqual(model.lineara.weight.device     ,			torch.device("""meta"""			)			)
										self.assertEqual(model.batchnorm.weight.device     ,			torch.device("""meta"""			)			)
										self.assertEqual(model.lineara.weight.device     ,			torch.device("""meta"""			)			)
										self.assertEqual(model.batchnorm.running_mean.device     ,			torch.device("""meta"""			)			)

										__magic_name__						:						Tuple										=torch.randn(2     ,			3			)
										__magic_name__						:						int										=model(__snake_case			)
										self.assertEqual(output.device     ,			__snake_case			)

										# Removing hooks loads back the weights in the model.
										remove_hook_from_module(model.lineara			)
										remove_hook_from_module(model.batchnorm			)
										remove_hook_from_module(model.lineara			)
										self.assertEqual(model.lineara.weight.device     ,			torch.device("""cpu"""			)			)
										self.assertEqual(model.batchnorm.weight.device     ,			torch.device("""cpu"""			)			)
										self.assertEqual(model.lineara.weight.device     ,			torch.device("""cpu"""			)			)

						def      A__    (		self :List[Any]			):
										'''simple docstring'''

										__magic_name__						:						Any										=ModelForTest()

										# Everything is on CPU
										self.assertEqual(model.lineara.weight.device     ,			torch.device("""cpu"""			)			)
										self.assertEqual(model.batchnorm.weight.device     ,			torch.device("""cpu"""			)			)
										self.assertEqual(model.lineara.weight.device     ,			torch.device("""cpu"""			)			)

										# This will move each submodule on different devices
										__magic_name__						:						str										=0 if torch.cuda.is_available() else """cpu"""
										attach_align_device_hook(__snake_case     ,			execution_device=__snake_case     ,			offload=__snake_case			)

										# Parameters have been offloaded, so on the meta device
										self.assertEqual(model.lineara.weight.device     ,			torch.device("""meta"""			)			)
										self.assertEqual(model.batchnorm.weight.device     ,			torch.device("""meta"""			)			)
										self.assertEqual(model.lineara.weight.device     ,			torch.device("""meta"""			)			)
										# Buffers are not included in the offload by default, so are on the execution device
										__magic_name__						:						Optional[Any]										=torch.device(__snake_case			)
										self.assertEqual(model.batchnorm.running_mean.device     ,			__snake_case			)

										__magic_name__						:						List[Any]										=torch.randn(2     ,			3			)
										__magic_name__						:						str										=model(__snake_case			)
										self.assertEqual(output.device     ,			__snake_case			)

										# Removing hooks loads back the weights in the model.
										remove_hook_from_submodules(__snake_case			)
										self.assertEqual(model.lineara.weight.device     ,			torch.device("""cpu"""			)			)
										self.assertEqual(model.batchnorm.weight.device     ,			torch.device("""cpu"""			)			)
										self.assertEqual(model.lineara.weight.device     ,			torch.device("""cpu"""			)			)

										# Now test with buffers included in the offload
										attach_align_device_hook(__snake_case     ,			execution_device=__snake_case     ,			offload=__snake_case     ,			offload_buffers=__snake_case			)

										# Parameters have been offloaded, so on the meta device, buffers included
										self.assertEqual(model.lineara.weight.device     ,			torch.device("""meta"""			)			)
										self.assertEqual(model.batchnorm.weight.device     ,			torch.device("""meta"""			)			)
										self.assertEqual(model.lineara.weight.device     ,			torch.device("""meta"""			)			)
										self.assertEqual(model.batchnorm.running_mean.device     ,			torch.device("""meta"""			)			)

										__magic_name__						:						Optional[int]										=torch.randn(2     ,			3			)
										__magic_name__						:						Union[str, Any]										=model(__snake_case			)
										self.assertEqual(output.device     ,			__snake_case			)

										# Removing hooks loads back the weights in the model.
										remove_hook_from_submodules(__snake_case			)
										self.assertEqual(model.lineara.weight.device     ,			torch.device("""cpu"""			)			)
										self.assertEqual(model.batchnorm.weight.device     ,			torch.device("""cpu"""			)			)
										self.assertEqual(model.lineara.weight.device     ,			torch.device("""cpu"""			)			)


						def      A__    (		self :List[str]			):
										'''simple docstring'''

										__magic_name__						:						Dict										=ModelForTest()

										# Everything is on CPU
										self.assertEqual(model.lineara.weight.device     ,			torch.device("""cpu"""			)			)
										self.assertEqual(model.batchnorm.weight.device     ,			torch.device("""cpu"""			)			)
										self.assertEqual(model.lineara.weight.device     ,			torch.device("""cpu"""			)			)

										# This will move each submodule on different devices
										__magic_name__						:						List[str]										=0 if torch.cuda.is_available() else """cpu"""
										attach_align_device_hook(
										    __snake_case     ,			execution_device=__snake_case     ,			offload=__snake_case     ,			weights_map=model.state_dict()			)

										# Parameters have been offloaded, so on the meta device
										self.assertEqual(model.lineara.weight.device     ,			torch.device("""meta"""			)			)
										self.assertEqual(model.batchnorm.weight.device     ,			torch.device("""meta"""			)			)
										self.assertEqual(model.lineara.weight.device     ,			torch.device("""meta"""			)			)
										# Buffers are not included in the offload by default, so are on the execution device
										__magic_name__						:						Optional[Any]										=torch.device(__snake_case			)
										self.assertEqual(model.batchnorm.running_mean.device     ,			__snake_case			)

										__magic_name__						:						int										=torch.randn(2     ,			3			)
										__magic_name__						:						Any										=model(__snake_case			)
										self.assertEqual(output.device     ,			__snake_case			)

										# Removing hooks loads back the weights in the model.
										remove_hook_from_submodules(__snake_case			)
										self.assertEqual(model.lineara.weight.device     ,			torch.device("""cpu"""			)			)
										self.assertEqual(model.batchnorm.weight.device     ,			torch.device("""cpu"""			)			)
										self.assertEqual(model.lineara.weight.device     ,			torch.device("""cpu"""			)			)

										# Now test with buffers included in the offload
										attach_align_device_hook(
										    __snake_case     ,			execution_device=__snake_case     ,			offload=__snake_case     ,			weights_map=model.state_dict()     ,			offload_buffers=__snake_case     ,			)

										# Parameters have been offloaded, so on the meta device, buffers included
										self.assertEqual(model.lineara.weight.device     ,			torch.device("""meta"""			)			)
										self.assertEqual(model.batchnorm.weight.device     ,			torch.device("""meta"""			)			)
										self.assertEqual(model.lineara.weight.device     ,			torch.device("""meta"""			)			)
										self.assertEqual(model.batchnorm.running_mean.device     ,			torch.device("""meta"""			)			)

										__magic_name__						:						List[Any]										=torch.randn(2     ,			3			)
										__magic_name__						:						str										=model(__snake_case			)
										self.assertEqual(output.device     ,			__snake_case			)

										# Removing hooks loads back the weights in the model.
										remove_hook_from_submodules(__snake_case			)
										self.assertEqual(model.lineara.weight.device     ,			torch.device("""cpu"""			)			)
										self.assertEqual(model.batchnorm.weight.device     ,			torch.device("""cpu"""			)			)
										self.assertEqual(model.lineara.weight.device     ,			torch.device("""cpu"""			)			)




