def test_weight_based_state_transfer(self):
    'Test that preproc layers can transfer state via get/set weights..'

    def get_model():
        input_data = keras.Input(shape=(1,))
        layer = AddingPreprocessingLayer()
        output = layer(input_data)
        model = keras.Model(input_data, output)
        model._run_eagerly = testing_utils.should_run_eagerly()
        return (model, layer)
    input_dataset = np.array([1, 2, 3, 4, 5])
    (model, layer) = get_model()
    layer.adapt(input_dataset)
    self.assertAllEqual([[16], [17], [18]], model.predict([1.0, 2.0, 3.0]))
    weights = model.get_weights()
    (model_2, _) = get_model()
    self.assertAllEqual([[1], [2], [3]], model_2.predict([1.0, 2.0, 3.0]))
    model_2.set_weights(weights)
    self.assertAllEqual([[16], [17], [18]], model_2.predict([1.0, 2.0, 3.0]))
Test that preproc layers can transfer state via get/set weights..
def test_weight_based_state_transfer(self):
    

    def get_model():
        input_data = keras.Input(shape=(1,))
        layer = AddingPreprocessingLayer()
        output = layer(input_data)
        model = keras.Model(input_data, output)
        model._run_eagerly = testing_utils.should_run_eagerly()
        return (model, layer)
    input_dataset = np.array([1, 2, 3, 4, 5])
    (model, layer) = get_model()
    layer.adapt(input_dataset)
    self.assertAllEqual([[16], [17], [18]], model.predict([1.0, 2.0, 3.0]))
    weights = model.get_weights()
    (model_2, _) = get_model()
    self.assertAllEqual([[1], [2], [3]], model_2.predict([1.0, 2.0, 3.0]))
    model_2.set_weights(weights)
    self.assertAllEqual([[16], [17], [18]], model_2.predict([1.0, 2.0, 3.0]))