title: When writing a custom loss function, should I use tf.reduce_mean, and if so how? Does it ever matter?
body: <p>The sample code below shows that all the following give the same (correct) results when
  writing a custom loss function (calculating mean_squared_error) for
  a simple linear regression model.</p>
  <ol>
  <li>Do not use tf_reduce_mean() (so returning a loss for each example)</li>
  <li>Use tf_reduce_mean() (so returning a single loss)</li>
  <li>Use tf_reduce_mean(..., axis-1)</li>
  </ol>
  <p>Is there any reason to prefer one approach to another, and are there any circumstances
  where it makes a difference?</p>
  <p>(There is, for example sample code at
  <a href="https://stackoverflow.com/questions/45961428/make-a-custom-loss-function-in-keras/">Make a custom loss function in keras</a>
  that suggests axis=-1 should be used)</p>
  <pre class="lang-py prettyprint-override"><code>import numpy as np
  import tensorflow as tf
  
  # Create simple dataset to do linear regression on
  # The mean squared error (~ best achievable MSE loss after fitting linear regression) for this dataset is 0.01
  xtrain = np.random.randn(5000)                  # Already normalized
  ytrain = xtrain + np.random.randn(5000) * 0.1   # Close enough to being normalized
  
  # Function to create model and fit linear regression, and report final loss
  def cre_and_fit(loss=&quot;mean_squared_error&quot;, lossdescription=&quot;&quot;,epochs=20):
      model = tf.keras.models.Sequential([tf.keras.layers.Dense(1, input_shape=(1,))]) 
      model.compile(loss=loss, optimizer=&quot;RMSProp&quot;)
      history = model.fit(xtrain, ytrain, epochs=epochs, verbose=False)
      print(f&quot;Final loss value for {lossdescription}: {history.history['loss'][-1]:.4f}&quot;)
  
  # Result from standard MSE loss ~ 0.01
  cre_and_fit(&quot;mean_squared_error&quot;,&quot;Keras standard MSE&quot;)
  
  # This gives the right result, not reducing. Return shape = (batch_size,)
  cre_and_fit(lambda y_true, y_pred: (y_true-y_pred)*(y_true-y_pred),
                 &quot;custom loss, not reducing over batch items&quot; )
  
  # This also gives the right result, reducing over batch items. Return shape = ()
  cre_and_fit(lambda y_true, y_pred: tf.reduce_mean((y_true-y_pred)*(y_true-y_pred) ),
                 &quot;custom loss, reducing over batch items&quot;)
  
  # How about using axis=-1? Also gives the same result
  cre_and_fit(lambda y_true, y_pred: tf.reduce_mean((y_true-y_pred)*(y_true-y_pred), axis=-1),
                 &quot;custom loss, reducing with axis=-1&quot; )
  </code></pre>