Does Keras support if statements inside the model?

Question

In my input_shape there is a variable x that goes from 0 to 100, and I want to separate the neural network such that for values smaller than 50 it uses some weights and for values bigger than 50 it uses other weights. I imagine something like this:

inputs = keras.Input(shape=(3,)) # let's say x is the first of the three variables

if inputs[0] < 50:
    x = layers.Dense(10)(inputs)
    output_small = layers.Dense(10)(x)
else:
    x = layers.Dense(10)(inputs)
    output_big = layers.Dense(10)(x)

model1 = keras.Model(inputs, output_small, name="small_values")
model2 = keras.Model(inputs, output_big, name="big_values")

I necessarily need the if inside the model because I later want to use output_small and output_big in a same Dense layer, therefore I can't run the two samples separately. Would this work or is it possible in another way using Keras? Otherwise is it possible using some other tool?

Answer 1

Yes you should be able to manage conditional branching and merging of layers in a model with functional API.

Link to guide on functional API: https://keras.io/guides/functional_api/