human
5. Batching
Fully vectorize the operations so that they work on a batch of images"
Please write code following the instructions in jupyter notebook style.

=====
Give me a code to 5?
assistant
#Here is code to 5.

class Grayscale(tf.keras.layers.Layer):
    def __init__(self, **kwargs):
        super(Grayscale, self).__init__(kwargs)
    
    def call(self, img):
        wt = tf.constant([[0.2126], [0.7152], [0.0722]]) # 3x1 matrix
        c_linear = tf.matmul(img, wt) # (N, ht,wd,3) x (3x1) -> (N, ht, wd)
        gray = tf.where(c_linear > 0.0031308,
                        1.055 * tf.pow(c_linear, 1/2.4) - 0.055,
                        12.92*c_linear)
        return gray # (N, ht, wd)

model = tf.keras.Sequential([
    Grayscale(input_shape=(336, 600, 3)),
    tf.keras.layers.Lambda(lambda gray: tf.reduce_mean(gray, axis=[1, 2])) # note axis change
])

%%time
ds = tf.data.TextLineDataset('input.txt').repeat(10).map(read_jpeg).batch(5)
overall = tf.constant([0.])
count = 0
for batch in ds:
    bm = model(batch)
    overall += tf.reduce_sum(bm)
    count += len(bm)
print(overall/count)