code: def add_weight(cls):
    
    @functools.wraps(cls.add_weight)
    def _add_weight(self,
                    name=None,
                    shape=None,
                    dtype=None,
                    initializer=None,
                    regularizer=None,
                    **kwargs):
      """Adds weight."""
      if isinstance(initializer, tf.keras.layers.Layer):
        weight = initializer(shape, dtype)
        self._trainable_weights.extend(initializer.trainable_weights)  # pylint: disable=protected-access
        self._non_trainable_weights.extend(initializer.non_trainable_weights)  # pylint: disable=protected-access
        if regularizer is not None:
          # TODO(trandustin): Replace need for this with
          # Layer._handle_weight_regularization. For Eager compatibility, random
          # variable __init__s cannot apply TF ops (cl/220898007).
          def loss_fn():
            """Creates a regularization loss `Tensor`."""
            with tf.name_scope(name + '/Regularizer'):
              return regularizer(initializer(shape, dtype))
          self.add_loss(loss_fn)
        return weight
      return super(cls, self).add_weight(name=name,
                                         shape=shape,
                                         dtype=dtype,
                                         initializer=initializer,
                                         regularizer=regularizer,
                                         **kwargs)
    cls.add_weight = _add_weight
    return cls
docs: Decorator for Layers, overriding add_weight for trainable initializers.