class_docstr: Distributed tests for clustering.
docstr: Verifies that associations of weights are updated per replica.
docstr_tok: ["Verifies","that","associations","of","weights","are","updated","per","replica","."]
code: def testAssociationValuesPerReplica(self, distribution):
      """Verifies that associations of weights are updated per replica."""
      assert tf.distribute.get_replica_context() is not None
      with distribution.scope():
          assert tf.distribute.get_replica_context() is None
          input_shape = (1, 2)
          output_shape = (2, 8)
          l = cluster_wrapper.ClusterWeights(
              keras.layers.Dense(8, input_shape=input_shape),
              number_of_clusters=2,
              cluster_centroids_init=CentroidInitialization.LINEAR,
          )
          l.build(input_shape)
  
          clusterable_weights = l.layer.get_clusterable_weights()
          self.assertLen(clusterable_weights, 1)
          weights_name = clusterable_weights[0][0]
          self.assertEqual(weights_name, "kernel")
          centroids1 = l.cluster_centroids[weights_name]
  
          mean_weight = tf.reduce_mean(l.layer.kernel)
          min_weight = tf.reduce_min(l.layer.kernel)
          max_weight = tf.reduce_max(l.layer.kernel)
          max_dist = max_weight - min_weight
  
          def assert_all_cluster_indices(per_replica, indices_val):
              if indices_val == 1:
                  val_tensor = tf.dtypes.cast(
                      tf.ones(shape=output_shape), per_replica[0].dtype
                  )
              if indices_val == 0:
                  val_tensor = tf.dtypes.cast(
                      tf.zeros(shape=output_shape), per_replica[0].dtype
                  )
              for i in range(0, len(per_replica)):
                  all_equal = tf.reduce_all(
                      tf.equal(per_replica[i], val_tensor)
                  )
                  self.assertTrue(all_equal)
  
          def update_fn(v, val):
              return v.assign(val)
  
          initial_val = tf.Variable(
              [mean_weight, mean_weight + 2.0 * max_dist],
              aggregation=tf.VariableAggregation.MEAN,
          )
  
          centroids1 = distribution.extended.update(
              centroids1, update_fn, args=(initial_val,)
          )
          l.call(tf.ones(shape=input_shape))
  
          clst_indices = l.pulling_indices[weights_name]
          per_replica = distribution.experimental_local_results(
              clst_indices
          )
          assert_all_cluster_indices(per_replica, 0)
  
          second_val = tf.Variable(
              [mean_weight - 2.0 * max_dist, mean_weight],
              aggregation=tf.VariableAggregation.MEAN,
          )
          centroids2 = l.cluster_centroids[weights_name]
          centroids2 = distribution.extended.update(
              centroids2, update_fn, args=(second_val,)
          )
          l.call(tf.ones(shape=input_shape))
  
          clst_indices = l.pulling_indices[weights_name]
          per_replica = distribution.experimental_local_results(
              clst_indices
          )
          assert_all_cluster_indices(per_replica, 1)
  
file_path: model-optimization/tensorflow_model_optimization/python/core/clustering/keras/cluster_distributed_test.py
imports: import itertools
  from absl.testing import parameterized
  import numpy
  import tensorflow
  from tensorflow_model_optimization.python.core.clustering.keras import cluster
  from tensorflow_model_optimization.python.core.clustering.keras import cluster_config
  from tensorflow_model_optimization.python.core.clustering.keras import cluster_wrapper
  from tensorflow_model_optimization.python.core.clustering.keras.experimental import cluster
  from tensorflow_model_optimization.python.core.keras import test_utils
label_desc: Use image node_id 2 for calling the ClusterDistributedTest obj's underlying member method code with example usage: obj.testAssociationValuesPerReplica(distribution) and returns: v