code: def test_update_mirrored_variable(self, strategy):
      # This is also uncommon since most model parameters should be updated by
      # optimizer, and this test case is for completeness.
      #
      # In the cross replica context, assigning to the variable assigns the same
      # value to all replicas. This is true with the loaded model as well.
      #
      # However in replica context, MirroredVariable (synchronization=ON_WRITE)
      # in a loaded model behaves differently. Updating MirroredVariable only
      # update the current replica's variable with the current replica's value.
      # There's no aggregation. This doesn't affect variables that are updated
      # through optimizer. This is work as intended but can be surprising.
  
      class Model(tf.Module):
          def __init__(self):
              self.v = tf.Variable(
                  0.0,
                  synchronization=tf.VariableSynchronization.ON_WRITE,
                  aggregation=tf.VariableAggregation.MEAN,
              )
  
          @tf.function(
              input_signature=[
                  tf.TensorSpec(shape=(), dtype=tf.float32)
              ]
          )
          def update(self, value):
              self.v.assign_add(value)
  
      export_dir = self.get_temp_dir()
      # value = strategy.experimental_distribute_values_from_function(
      #     lambda ctx: tf.identity([1., 2.][ctx.replica_id_in_sync_group]))
      with strategy.scope():
          m = Model()
          tf.saved_model.save(m, export_dir)
          del m
  
      with strategy.scope():
          loaded = tf.saved_model.load(export_dir)
      self.assertAllEqual(
          self.evaluate(strategy.experimental_local_results(loaded.v)),
          [0.0, 0.0],
      )
      self.evaluate(loaded.v.assign(1.0))
      self.assertAllEqual(
          self.evaluate(strategy.experimental_local_results(loaded.v)),
          [1.0, 1.0],
      )
  
imports: import os
  from absl.testing import parameterized
  import tensorflow
  import tensorflow.compat.v1
  from tensorflow.python.distribute import combinations
  from tensorflow.python.distribute import multi_worker_test_base
  from tensorflow.python.distribute import parameter_server_strategy_v2
  from tensorflow.python.distribute import sharded_variable
  from tensorflow.python.distribute import strategy_combinations
  from tensorflow.python.distribute import test_util
  from tensorflow.python.distribute import values
  from tensorflow.python.eager import context
  from tensorflow.python.eager import test
  from tensorflow.python.ops import lookup_ops
label_desc: Use image node_id 4 for calling the SaveAndLoadForTrainingTest obj's underlying member method code with example usage: obj.test_update_mirrored_variable(strategy) without return types