code: def test_train_with_distribution_strategy(self):
      (
          w,
          x0,
          y0,
          lr,
          adv_config,
          w_new,
      ) = self._set_up_linear_regression()
      inputs = tf.data.Dataset.from_tensor_slices(
          {"feature": x0, "label": y0}
      ).batch(NUM_REPLICAS)
  
      strategy = self._get_mirrored_strategy()
      with strategy.scope():
          # Makes sure we are running on multiple devices.
          self.assertEqual(NUM_REPLICAS, strategy.num_replicas_in_sync)
          model = build_linear_keras_functional_model(
              input_shape=(2,), weights=w
          )
          adv_model = (
              adversarial_regularization.AdversarialRegularization(
                  model, label_keys=["label"], adv_config=adv_config
              )
          )
          adv_model.compile(
              optimizer=tf.keras.optimizers.SGD(lr), loss="MSE"
          )
  
      adv_model.fit(x=inputs)
  
      # The updated weight should be the same regardless of the number of devices.
      self.assertAllClose(
          w_new, tf.keras.backend.get_value(model.weights[0])
      )
  
file_path: neural-structured-learning/neural_structured_learning/keras/adversarial_regularization_multi_device_test.py
imports: from __future__ import absolute_import
  from __future__ import division
  from __future__ import print_function
  import neural_structured_learning.configs
  from neural_structured_learning.keras import adversarial_regularization
  import numpy
  import tensorflow
label_desc: Use image node_id 4 for calling the AdversarialRegularizationMultiDeviceTest obj's underlying member method code with example usage: obj.test_train_with_distribution_strategy() without return types