examples/demo_jax_distributed.py/make_model
def make_model():
    input = keras.Input(shape=[28, 28, 1])
    y = make_backbone()(input)
    y = keras.layers.Flatten()(y)
    y = keras.layers.Dense(200, activation="relu")(y)
    y = keras.layers.Dropout(0.4)(y)
    y = keras.layers.Dense(10, activation="softmax")(y)
    model = keras.Model(inputs=input, outputs=y)
    return model