code: def test_train_with_distribution_strategy(self, model_fn):
      (
          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(1)
  
      strategy = tf.distribute.MirroredStrategy()
      with strategy.scope():
          model = model_fn(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",
              metrics=["mae"],
          )
  
      adv_model.fit(x=inputs)
  
      self.assertAllClose(
          w_new, tf.keras.backend.get_value(model.weights[0])
      )
  
imports: from __future__ import absolute_import
  from __future__ import division
  from __future__ import print_function
  import collections
  import os
  from absl.testing import parameterized
  import neural_structured_learning.configs
  from neural_structured_learning.keras import adversarial_regularization
  import numpy
  import tensorflow
  from tensorflow.python.framework import test_util
label_desc: Use image node_id 7 for calling the AdversarialRegularizationTest obj's underlying member method code with example usage: obj.test_train_with_distribution_strategy(model_fn) without return types