Can't get Keras functional model to work: "could not broadcast input array..." For a really trivial MWE, here is some code that I expected to work and does not:

import tensorflow.keras as keras
import numpy as np
input_ = keras.layers.Input((256, 768))
layer = keras.layers.Lambda(lambda x: x[0, :], output_shape=(768,),    
                            input_shape=(256, 768))
model = keras.models.Model(inputs = [input_], outputs=layer(input_))
model.predict(np.zeros((1,256,768)))


I expect this to return np.zeros((1, 768)), but instead it throws the following exception:

Traceback (most recent call last):
  File "&lt;stdin&gt;", line 1, in &lt;module&gt;
  File "/Users/holmes5/anaconda3/lib/python3.7/site-packages/tensorflow/python/keras/engine/training.py", line 1078, in predict
callbacks=callbacks)
  File "/Users/holmes5/anaconda3/lib/python3.7/site-packages/tensorflow/python/keras/engine/training_arrays.py", line 370, in 
model_iteration
aggregator.aggregate(batch_outs, batch_start, batch_end)
  File "/Users/holmes5/anaconda3/lib/python3.7/site-packages/tensorflow/python/keras/engine/training_utils.py", line 169, in aggregate
self.results[i][batch_start:batch_end] = batch_out
ValueError: could not broadcast input array from shape (256,768) into shape (1,768)


What am I doing wrong?
Fixed it, sorry. I just needed to say lambda x: x[:, 0, :], so it included the sample index.