How to use only the latest element of a sequence in the right branch of this nn?

hi, i am having issue with the input shape. is there any way to force tensorflow to use only the last element of a sequence for the mlp branch of this neural network?

mlp_two = tf.keras.models.Sequential(name="mlp")

mlp_two.add(tf.keras.layers.Lambda(lambda x: x[-1] ))

mlp_two.add(ConcreteDenseDropout(Dense(50,activation= "relu"), is_mc_dropout=True, weight_regularizer=wr, dropout_regularizer=dr,temperature=0.1)) 

mlp_two.add(ConcreteDenseDropout(Dense(50,activation= "relu"), is_mc_dropout=True, weight_regularizer=wr, dropout_regularizer=dr,temperature=0.1))

mlp_two.add(Dense(1,activation= "softplus"))
rnn_model_two = tf.keras.models.Sequential(name="lstm")

rnn_model_two.add(tf.keras.layers.LSTM(units=100, return_sequences=True, activation='relu',recurrent_dropout=0.2, input_shape=[3,1]))

rnn_model_two.add(tf.keras.layers.LSTM(units=50, activation='relu',recurrent_dropout=0.2))

rnn_model_two.add( Dense(100,activation= "relu") )

rnn_model_two.add(tf.keras.layers.Dense(1))
def normal_sp(params): 

        #return tfd.Normal(loc=params[:,0:1], scale=1e-5 +  params[:,1:2]) 

        return tfd.Normal(loc=params[:,0:1], scale=1e-3 + tf.math.softplus(0.005*  params[:,1:2]))
def fun_two():

  inputs = Input(shape=(3,1),name="input layer")

  first=rnn_model_two(inputs,training=True)

  second=mlp_two(inputs,training=True)

  #first=tf.keras.layers.Flatten()(first)

  #second=tf.keras.layers.Flatten()(second)

  z=tf.concat([first, second  ],1)

  #z=tf.keras.layers.Flatten()(z)

  dist_mc = tfp.layers.DistributionLambda(normal_sp, name='normal_sp')(z) 

  return Model(inputs=inputs, outputs=dist_mc)

sliding_BNN = fun_two()
callback = tf.keras.callbacks.EarlyStopping(monitor=‘val_loss’, patience=2000)

optimizer = tf.optimizers.SGD(learning_rate=0.0001,momentum=0.9)

sliding_BNN.compile(optimizer=optimizer,

              loss=NLL  ,metrics= [tf.keras.metrics.RootMeanSquaredError()]

             ) 

sliding_BNN.build(input_shape=(3,1))

sliding_BNN.summary()

tf.keras.utils.plot_model(sliding_BNN, “rnn_sliding.png”, show_shapes=True)

basically i wish to use only the latest element of the time sequence as input of the mlp because it’s scope is to predict the variance.
this means that in tf.concat_23 should receive 2 scalars that then are passed to normal_sp

1 Like

Hello @P11

Thank you for using TensorFlow,
In the code, your current Lambda layer is selects the last element along the 1st axis, which is the batch axis.
To fix this, you need to modify the Lambda layer in your mlp_two
mlp_two.add(tf.keras.layers.Lambda(lambda x: x[:, -1, :]))

Thank you