MaxPooling 1D layer in Trax

Question

How can I code my own MaxPooling_1D layer in google Trax? I understand that current max pooling is based on 2D max pooling.

Here's how I tried using Keras 1d layer

   import trax.layers as tl
   def computePool(max_pool_1d,in_tensor):
      #print(in_tensor)
      return max_pool_1d(in_tensor)
   def maxPooling1D():
      max_pool_1d = tf.keras.layers.GlobalMaxPooling1D()
      return tl.Fn('maxPooling1D', lambda x: computePool(max_pool_1d, x))

But this wouldn't go through in the model. I want to create a layer with pool size = 2

Answer 1

The answer is

tl.MaxPool(pool_size=(2,), strides=(1,), padding='SAME'),