def METHOD_NAME(df_example):
    df = df_example
    df = df[:1_000]  # To make the tests run faster
    df_train, df_valid, df_test = df.split_random([0.8, 0.1, 0.1], random_state=42)
    features = ['vx', 'vy', 'vz']
    target = 'FeH'
    # Scaling the features
    df_train = df_train.ml.minmax_scaler(features=features)
    features = df_train.get_column_names(regex='^minmax_')
    # Apply preprocessing to the validation
    state_prep = df_train.state_get()
    df_valid.state_set(state_prep)
    train_gen = df_train.ml.tensorflow.to_keras_generator(features=features, target=target, batch_size=128)
    valid_gen = df_valid.ml.tensorflow.to_keras_generator(features=features, target=target, batch_size=128)
    # Create the model
    nn_model = K.Sequential()
    nn_model.add(K.layers.Dense(3, activation='tanh'))
    nn_model.add(K.layers.Dense(1, activation='linear'))
    nn_model.compile(optimizer='sgd', loss='mse', metrics=[r2_keras])
    nn_model.fit(x=train_gen, validation_data=valid_gen, epochs=5, steps_per_epoch=7, validation_steps=1, verbose=0)
    keras_model = vaex.ml.tensorflow.KerasModel(features=features, prediction_name='pred', model=nn_model, custom_objects={'r2_keras': r2_keras})
    df_train = keras_model.transform(df_train)
    # The final state transfer
    state_final = df_train.state_get()
    df_valid.state_set(state_final)
    df_test.state_set(state_final)
    assert 'pred' in df_train
    assert 'pred' in df_valid
    assert 'pred' in df_test
    assert df_valid.shape == (100, 15)
    assert df_test.shape == (100, 15)