code: def test_fit_simple_linear_model():
      seed = 0x2019
      np.random.seed(seed)
      tf.random.set_seed(seed)
      num_examples = 100000
      x = np.random.standard_normal((num_examples, 3))
      w = np.random.standard_normal((3, 1))
      y = (
          np.dot(x, w)
          + np.random.standard_normal((num_examples, 1)) * 1e-4
      )
  
      model = tf.keras.models.Sequential()
      model.add(tf.keras.layers.Dense(input_shape=(3,), units=1))
      # using num_examples - 1 since steps starts from 0.
      optimizer = SWA(
          "sgd",
          start_averaging=num_examples // 32 - 1,
          average_period=100,
      )
      model.compile(optimizer, loss="mse")
      model.fit(x, y, epochs=2)
      optimizer.assign_average_vars(model.variables)
  
      x = np.random.standard_normal((100, 3))
      y = np.dot(x, w)
  
      predicted = model.predict(x)
  
      max_abs_diff = np.max(np.abs(predicted - y))
      assert max_abs_diff < 1e-3
  
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 6 for calling a global function with example usage: test_fit_simple_linear_model() without return types