deepkoopman_second_input.py

# Add onto our existing x0/code0 computation graph
# x1 is the input x_{k+1}
x1 = tf.keras.Input(shape=INPUT_DIM, name='x1')
# we can re-use the same sub-models (encoder, decoder, linear_dynamics)
# to share the same weights
code1 = encoder(x1)
x1_pred = decoder(code1_pred)

model = tf.keras.Model(
    inputs={'x0': x0, 'x1': x1},
    outputs={'x0_reconstructed': x0_reconstructed,
             'y1': code1, 'y1_pred': code1_pred,
             'x1_pred': x1_pred}
)