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)

image

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