Remove the extra 2 dimensions in the output of TF reduce_window.

Author: DarrenZhang01

tf_helpers/lax.py

@@ -311,7 +311,7 @@ def reduce_window(inputs, init_value, reducer, window_dimensions, strides,
   # input shape is already finished in apply_fun of Pooling in stax.
   pooling = "AVG" if reducer == np.add else "MAX"
   output = nn.pool(inputs, window_dimensions, pooling, strides, padding)
-  return np.asarray(output)
+  return np.squeeze(np.asarray(output), axis=(0, output.ndim - 1))
 
 
 # TOTO (Zhibo Zhang): Expand the test cases of general convolution and revise