code: def test_subclass_model():
      class Block(tf.keras.Model):
          def __init__(self, units):
              super().__init__()
              self.dense1 = tf.keras.layers.Dense(units)
              self.dense2 = tf.keras.layers.Dense(units)
  
          def call(self, x):
              return self.dense2(self.dense1(x))
  
      class Custom(tf.keras.Model):
          def __init__(self):
              super().__init__()
              self.block1 = Block(16)
              self.block2 = Block(32)
  
          def call(self, x):
              return self.block2(self.block1(x))
  
      model = Custom()
      model.build(input_shape=(None, 4))
  
      optimizers_and_layers = [
          (tf.keras.optimizers.SGD(learning_rate=0.0), model.block1),
          (tf.keras.optimizers.Adam(), model.block2),
      ]
  
      block1_weights_before_train = [
          weight.numpy() for weight in model.block1.weights
      ]
      block2_weights_before_train = [
          weight.numpy() for weight in model.block2.weights
      ]
  
      multi_optimizer = MultiOptimizer(optimizers_and_layers)
  
      x = np.ones((128, 4)).astype(np.float32)
      y = np.ones((128, 32)).astype(np.float32)
      mse = tf.keras.losses.MeanSquaredError()
  
      for _ in range(10):
          for i in range(0, 128, 32):
              x_batch = x[i : i + 32]
              y_batch = y[i : i + 32]
              with tf.GradientTape() as tape:
                  loss = mse(y_batch, model(x_batch))
  
              grads = tape.gradient(loss, model.trainable_variables)
              multi_optimizer.apply_gradients(
                  zip(grads, model.trainable_variables)
              )
  
      block1_weights_after_train = [
          weight.numpy() for weight in model.block1.weights
      ]
      block2_weights_after_train = [
          weight.numpy() for weight in model.block2.weights
      ]
  
      assert_list_allclose(
          block1_weights_before_train, block1_weights_after_train
      )
      assert_list_not_allclose(
          block2_weights_before_train, block2_weights_after_train
      )
      assert_sync_iterations(
          multi_optimizer, desired=ceil(128 / 32) * 10
      )
  
file_path: addons/tensorflow_addons/optimizers/tests/discriminative_layer_training_test.py
imports: from math import ceil
  from packaging.version import Version
  import pytest
  import numpy
  import tensorflow
  from tensorflow_addons.optimizers.discriminative_layer_training import MultiOptimizer
  from tensorflow_addons.utils import test_utils
label_desc: Use image node_id 7 for calling a global function with example usage: test_subclass_model() and returns: self, self