I am building a federated learning model using Tensorflow federated, and I am following the tutorials provided in the official documentation. As I can see, most of the implementations provided are using a neural network as the local ML model. As I just did in the following snippet.

def create_keras_model():
  initializer = tf.keras.initializers.Zeros()
  return tf.keras.models.Sequential([
      tf.keras.layers.Input(shape=(9,)),
      tf.keras.layers.Dense(4, kernel_initializer=initializer),
      tf.keras.layers.Softmax(),
  ])
def model_fn():
  keras_model = create_keras_model()
  return tff.learning.from_keras_model(
      keras_model,
      input_spec=train_data[0].element_spec,
      loss=tf.keras.losses.SparseCategoricalCrossentropy(),
      metrics=[tf.keras.metrics.SparseCategoricalAccuracy()])

Since I am building a multi-classification model with (9) features and (4) target labels. Can I use a different ML model for local training, like (logistic regression )? and how can I adjust that?