code: def test_averaging():
      start_averaging = 0
      average_period = 1
      if hasattr(tf.keras.optimizers, "legacy"):
          sgd = tf.keras.optimizers.legacy.SGD(lr=1.0)
      else:
          sgd = tf.keras.optimizers.SGD(lr=1.0)
      optimizer = SWA(sgd, start_averaging, average_period)
  
      val_0 = [1.0, 1.0]
      val_1 = [2.0, 2.0]
      var_0 = tf.Variable(val_0)
      var_1 = tf.Variable(val_1)
  
      grad_val_0 = [0.1, 0.1]
      grad_val_1 = [0.1, 0.1]
      grad_0 = tf.constant(grad_val_0)
      grad_1 = tf.constant(grad_val_1)
      grads_and_vars = list(zip([grad_0, grad_1], [var_0, var_1]))
  
      optimizer.apply_gradients(grads_and_vars)
      optimizer.apply_gradients(grads_and_vars)
      optimizer.apply_gradients(grads_and_vars)
  
      np.testing.assert_allclose(
          var_1.read_value(), [1.7, 1.7], rtol=1e-06, atol=1e-06
      )
      np.testing.assert_allclose(
          var_0.read_value(), [0.7, 0.7], rtol=1e-06, atol=1e-06
      )
  
      optimizer.assign_average_vars([var_0, var_1])
  
      np.testing.assert_allclose(var_0.read_value(), [0.8, 0.8])
      np.testing.assert_allclose(var_1.read_value(), [1.8, 1.8])
  
file_path: addons/tensorflow_addons/optimizers/tests/stochastic_weight_averaging_test.py
imports: import numpy
  import pytest
  import tensorflow
  from tensorflow_addons.optimizers import stochastic_weight_averaging
  from tensorflow_addons.optimizers.utils import fit_bn
label_desc: Use image node_id 1 for calling a global function with example usage: test_averaging() without return types