I managed to get this working now and there was a similar problem on SO here 69.
I’ll post the working code here in case anyone is interested.
import tensorflow as tf

# create a new layer that simple multiplies the input by a scalar
# got this from
# https://www.tensorflow.org/api_docs/python/tf/keras/layers/Lambda
class ScaleLayer(tf.keras.layers.Layer):
    def __init__(self, scale):
        super(ScaleLayer, self).__init__()
        self.scale = scale
    def call(self, inputs):
        return inputs * self.scale


# need a function to create sub-networks
def build_subnetwork(input_tensor, units, activation='relu', n_hidden=1):
    x = tf.keras.layers.Dense(units, activation=activation)(input_tensor)
    for _ in range(1, n_hidden):
        x = tf.keras.layers.Dense(units, activation=activation)(x)

    # single output here?
    x = tf.keras.layers.Dense(1, activation='linear')(x)

    return x

def build_model(
    input_shape = 1,
    output_shape = 1,
    units = 128,
    activation = 'relu',
    n_hidden_subnetwork = 2,
    scales = [1]
    ):
    input_layer = tf.keras.Input(shape=(input_shape,))
    # create sub-networks
    xs = []
    for scale in scales:
        scaled_input = ScaleLayer(scale)(input_layer)
        xs.append(build_subnetwork(scaled_input, units, activation, n_hidden_subnetwork))
    if len(xs) > 1:
        output_layer = tf.keras.layers.add([x for x in xs])
    else:
        output_layer = xs[0]
    output_layer = tf.keras.layers.Dense(output_shape, activation="linear")(output_layer)
    model = tf.keras.models.Model(inputs=input_layer, outputs=output_layer)
    return model