message: Document scaling behavior of `compute_gradients()`
diff: @@ -111,12 +111,22 @@ class CrossShardOptimizer(optimizer.Optimizer):
  def compute_gradients(self, loss, var_list=None, **kwargs):
  """Compute gradients of "loss" for the variables in "var_list".
  - This simply wraps the compute_gradients() from the real optimizer. The
  - gradients will be aggregated in the apply_gradients() so that user can
  + This simply wraps `compute_gradients()` from the real optimizer. The
  + gradients will be aggregated in `apply_gradients()` so that user can
  modify the gradients like clipping with per replica global norm if needed.
  The global norm with aggregated gradients can be bad as one replica's huge
  gradients can hurt the gradients from other replicas.
  + When the CrossShardOptimizer is constructed with
  + `reduction == losses.Reduction.MEAN` (default), this function scales the
  + loss by `1.0 / num_shards` before computing the gradients. Assuming the
  + optimizer uses the default implementation of `compute_gradients()`, the
  + gradients of the scaled loss are scaled by `1.0 / num_shards` compared to
  + the gradients of the original loss. This scaling factor is important because
  + `apply_gradients()` sums gradients across shards, rather than averaging
  + them. However, the scaling factor must be taken into account when clipping
  + the norm of the gradients or performing other postprocessing.
  +
  Args:
  loss: A Tensor containing the value to minimize.
  var_list: Optional list or tuple of `tf.Variable` to update to minimize
  