title: sum/avg/mean ops of tf.keras.backend do not work correctly over batch axis
body: <p>When adding an aggregation function on the batch axis (axis=0), everything is fine, as long as the batchsize is &lt;=32, with batchsize &gt;32 the output shape changes.</p>
  <pre><code>def model_test(samples, parameters):
  
      input_shape = (parameters)
      in_ = keras.layers.Input(input_shape, batch_size=samples)
  
      l1_size1 = parameters//2
      l2_size2 = parameters//3
  
      l1 = keras.layers.Dense(l1_size1,activation='relu')(in_)
      l2 = keras.layers.Dense(l2_size2,activation='relu')(l1)
  
      out_max = keras.backend.max(l2,axis=0)
  
      model = keras.models.Model(in_, out_max)
  
      return model
  </code></pre>
  <p>creating the model for 32 samples:</p>
  <pre><code>Model: &quot;functional_246&quot;
  _________________________________________________________________
  Layer (type)                 Output Shape              Param #   
  =================================================================
  input_146 (InputLayer)       [(32, 10)]                0         
  _________________________________________________________________
  dense_66 (Dense)             (32, 5)                   55        
  _________________________________________________________________
  dense_67 (Dense)             (32, 3)                   18        
  _________________________________________________________________
  tf_op_layer_Max_138 (TensorF [(3,)]                    0         
  =================================================================
  Total params: 73
  Trainable params: 73
  Non-trainable params: 0
  </code></pre>
  <p>running the model and display the real output shape:</p>
  <pre><code>x = tf.random.normal((orders ,parameters))
  result = model.predict(x)
  print(result.shape)
  
  
  (3,)
  </code></pre>
  <p>creating it for 33 samples, creates a similar model, with the same output and the same number of parameters, as expected:</p>
  <pre><code>orders = 33
  parameters = 10
  model = model_test(orders, parameters)
  model.summary()
  
  Model: &quot;functional_248&quot;
  _________________________________________________________________
  Layer (type)                 Output Shape              Param #   
  =================================================================
  input_147 (InputLayer)       [(33, 10)]                0         
  _________________________________________________________________
  dense_68 (Dense)             (33, 5)                   55        
  _________________________________________________________________
  dense_69 (Dense)             (33, 3)                   18        
  _________________________________________________________________
  tf_op_layer_Max_139 (TensorF [(3,)]                    0         
  =================================================================
  Total params: 73
  Trainable params: 73
  Non-trainable params: 0
  _________________________________________________________________
  </code></pre>
  <p>but the real output shape doubles:</p>
  <pre><code>x = tf.random.normal((orders ,parameters))
  result = model.predict(x)
  print(result.shape)
  
  (6,)
  </code></pre>
  <p>The observed output is always a multiple of batchsize//32 + 1 and hence does not match a actual expected output.</p>
  <p>This was tested with tensorflow 2.3.1 same on GPU and CPU.</p>
  <p>adding a further dimension, gives the same output, so this really only happens on axis=0</p>
  <pre><code>def model_test(samples, parameters):
  
      input_shape = (samples, parameters)
      in_ = keras.layers.Input(input_shape, batch_size=None)
  
      l1_size1 = parameters//2
      l2_size2 = parameters//3
  
      l1 = keras.layers.Dense(l1_size1,activation='relu')(in_)
      l2 = keras.layers.Dense(l2_size2,activation='relu')(l1)
  
      out_max = keras.backend.max(l2,axis=1)
  
      model = keras.models.Model(in_, out_max)
  
      return model
  </code></pre>