Is there a good method to reshape for the MaxPooling1D?

Question

I'm trying to use MaxPooling1D on "b". Then concatenate a and b.

But GlobalAveragePooling2D output is 2D tensor, MaxPooling1D input/output are 3D tensor.

So I use custom layer to reshape b for MaxPooling1D and Concatenate.

Is there any good method to replace these two custom layer?

import tensorflow as tf

a= tf.keras.layers.GlobalAveragePooling2D()(a)
b= tf.keras.layers.GlobalAveragePooling2D()(b)

#custom layer  reshape b from 2D tensor to 3D tensor for Maxpooling
b= tf.keras.layers.MaxPooling1D(pool_size=2,strides=2, padding='valid')(b)
#custom layer  reshape b from 3D tensor to 2D tensor for concatenate

AB = tf.keras.layers.Concatenate()([a, b])

Answer 1

Something like this would work:

import numpy as np
import tensorflow as tf

x_a = np.random.rand(100, 28, 28, 3)
x_b = np.random.rand(100, 28, 28, 3)

a_input = tf.keras.Input(shape=(28, 28, 3))
b_input = tf.keras.Input(shape=(28, 28, 3))

a = tf.keras.layers.GlobalAveragePooling2D()(a_input)
b = tf.keras.layers.GlobalAveragePooling2D()(b_input)

b = tf.keras.layers.Reshape((3, 1))(b)
b = tf.keras.layers.MaxPooling1D(pool_size=2, strides=2, padding='valid')(b)
b = tf.keras.layers.Flatten()(b)

AB = tf.keras.layers.Concatenate()([a, b])

model = tf.keras.Model([a_input, b_input], AB)

model.summary()