code: def _keras_build_fn(lr=0.01):
      layers = [
          Dense(512, input_shape=(784,), activation="relu"),
          Dense(10, input_shape=(512,), activation="softmax"),
      ]
  
      model = Sequential(layers)
  
      opt = tf.keras.optimizers.SGD(learning_rate=lr)
      model.compile(
          loss="categorical_crossentropy",
          optimizer=opt,
          metrics=["accuracy"],
      )
      return model
  
file_path: dask-ml/tests/model_selection/test_keras.py
imports: import numpy
  import pandas
  import pytest
  from distributed import Nanny
  from distributed.utils_test import gen_cluster
  from packaging import version
  from scipy.stats import loguniform
  from sklearn.datasets import make_classification
  from dask_ml.model_selection import IncrementalSearchCV
label_desc: Use image node_id 1 for calling a global function with example usage: _keras_build_fn(lr) and returns: model