message: Explicitly mention that the mean is taken over the last dimension.
diff: @@ -823,7 +823,10 @@ class Huber(LossFunctionWrapper):
  def mean_squared_error(y_true, y_pred):
  """Computes the mean squared error between labels and predictions.
  - `loss = square(y_true - y_pred)`
  + After computing the squared distance between the inputs, the mean value over
  + the last dimension is returned.
  +
  + `loss = mean(square(y_true - y_pred), axis=-1)`
  Args:
  y_true: Ground truth values. shape = `[batch_size, d0, .. dN]`.
  