code: def testStrategyExtendedUpdate(
      self, distribution, synchronization, aggregation
  ):
      if len(distribution.extended.parameter_devices) != 2:
          self.skipTest("n/a: needs exactly two parameter devices")
      if (
          synchronization
          == variables_lib.VariableSynchronization.ON_WRITE
          and aggregation != variables_lib.VariableAggregation.NONE
      ):
          self.skipTest(
              "n/a: doesn't apply to ON_WRITE variable with aggregation"
          )
      with distribution.scope():
          v = variables_lib.Variable(
              0.0,
              synchronization=synchronization,
              aggregation=aggregation,
          )
      value = values_lib.PerReplica([1.0, 2.0])
  
      assign_fn = lambda var, value: var.assign(value)
      self.evaluate(
          distribution.extended.update(v, assign_fn, args=(value,))
      )
      self.assertAllEqual(self.evaluate(v.values), [1.0, 2.0])
  
      assign_add_fn = lambda var, value: var.assign_add(value)
      self.evaluate(
          distribution.extended.update(v, assign_add_fn, args=(value,))
      )
      self.assertAllEqual(self.evaluate(v.values), [2.0, 4.0])
  
      assign_sub_fn = lambda var, value: var.assign_sub(value)
      self.evaluate(
          distribution.extended.update(v, assign_sub_fn, args=(value,))
      )
      self.assertAllEqual(self.evaluate(v.values), [1.0, 2.0])
  
      read_assign_fn = lambda var, value: var.assign_add(
          var.value() + var.read_value()
      )
      self.evaluate(
          distribution.extended.update(v, read_assign_fn, args=(value,))
      )
      self.assertAllEqual(self.evaluate(v.values), [3.0, 6.0])
  
imports: import copy
  import os
  from absl.testing import parameterized
  from tensorflow.python.checkpoint import checkpoint
  from tensorflow.python.checkpoint import checkpoint_options
  from tensorflow.python.distribute import collective_all_reduce_strategy
  from tensorflow.python.distribute import combinations
  from tensorflow.python.distribute import distribute_lib
  from tensorflow.python.distribute import distribute_utils
  from tensorflow.python.distribute import packed_distributed_variable
  from tensorflow.python.distribute import parameter_server_strategy
  from tensorflow.python.distribute import ps_values
  from tensorflow.python.distribute import strategy_combinations
  from tensorflow.python.distribute import test_util
  from tensorflow.python.distribute import tpu_strategy
  from tensorflow.python.distribute import values
  from tensorflow.python.eager import context
  from tensorflow.python.eager import def_function
  from tensorflow.python.eager import test
  from tensorflow.python.framework import constant_op
  from tensorflow.python.framework import dtypes
  from tensorflow.python.framework import indexed_slices
  from tensorflow.python.framework import ops
  from tensorflow.python.framework import tensor_shape
  from tensorflow.python.ops import array_ops
  from tensorflow.python.ops import check_ops
  from tensorflow.python.ops import control_flow_assert
  from tensorflow.python.ops import math_ops
  from tensorflow.python.ops import variable_scope
  from tensorflow.python.ops import variables
  from tensorflow.python.saved_model import save
  from tensorflow.python.saved_model import save_context
  from tensorflow.python.saved_model import save_options
  from tensorflow.python.types import core
label_desc: Use image node_id 10 for calling the DistributedVariableTest obj's underlying member method code with example usage: obj.testStrategyExtendedUpdate(distribution, synchronization, aggregation) without return types