class_docstr: An optimizer that averages gradients across TPU shards.
docstr: Construct a new cross-shard optimizer.
  
  Args:
    opt: An existing `Optimizer` to encapsulate.
    reduction: The reduction to apply to the shard losses.
    name: Optional name prefix for the operations created when applying
      gradients. Defaults to "CrossShardOptimizer".
    group_assignment: Optional 2d int32 lists with shape
      [num_groups, num_replicas_per_group] which describles how to apply
      optimizer to subgroups.
  
  Raises:
    ValueError: If reduction is not a valid cross-shard reduction.
docstr_tok: ["Construct","a","new","cross-shard","optimizer",".","Args",":","opt",":","An","existing","`","Optimizer","`","to","encapsulate",".","reduction",":","The","reduction","to","apply","to","the","shard","losses",".","name",":","Optional","name","prefix","for","the","operations","created","when","applying","gradients",".","Defaults","to","``","CrossShardOptimizer","''",".","group_assignment",":","Optional","2d","int32","lists","with","shape","[","num_groups",",","num_replicas_per_group","]","which","describles","how","to","apply","optimizer","to","subgroups",".","Raises",":","ValueError",":","If","reduction","is","not","a","valid","cross-shard","reduction","."]
code: def __init__(
      self,
      opt,
      reduction=losses.Reduction.MEAN,
      name="CrossShardOptimizer",
      group_assignment=None,
  ):
      """Construct a new cross-shard optimizer.
  
      Args:
        opt: An existing `Optimizer` to encapsulate.
        reduction: The reduction to apply to the shard losses.
        name: Optional name prefix for the operations created when applying
          gradients. Defaults to "CrossShardOptimizer".
        group_assignment: Optional 2d int32 lists with shape
          [num_groups, num_replicas_per_group] which describles how to apply
          optimizer to subgroups.
  
      Raises:
        ValueError: If reduction is not a valid cross-shard reduction.
      """
      accepted_reductions = (
          losses.Reduction.SUM,
          losses.Reduction.MEAN,
      )
      if reduction not in accepted_reductions:
          raise ValueError(
              f"Argument `reduction` should be one of {accepted_reductions}. "
              f"Received: {reduction}"
          )
      if not isinstance(opt, optimizer.Optimizer):
          raise TypeError(
              "CrossShardOptimizer only works with tf.training.Optimizer and not "
              f"Keras Optimizer. Received: {opt}. "
              "If you are using TPUStrategy, "
              "Keras Optimizer will sum gradients across replicas."
              "If you are using TPUEstimator, you may instead sum your gradients "
              "with:\n"
              "`grads = [tf.compat.v1.tpu.cross_replica_sum(g) for g in grads]`\n"
              "If you want to average your gradients, rescale your loss with: "
              "`loss /= global_batch_size`"
          )
  
      super(CrossShardOptimizer, self).__init__(False, name)
      self._opt = opt
      self._reduction = reduction
      self._group_assignment = group_assignment
  
file_path: tensorflow-upstream/tensorflow/python/tpu/tpu_optimizer.py
imports: from tensorflow.python.framework import ops
  from tensorflow.python.ops.losses import losses
  from tensorflow.python.platform import tf_logging
  from tensorflow.python.tpu import tpu_function
  from tensorflow.python.tpu.ops import tpu_ops
  from tensorflow.python.training import optimizer
  from tensorflow.python.util.tf_export import tf_export
label_desc: Use image node_id 1 to create a new CrossShardOptimizer object from inherited base classes: optimizer with example: obj = CrossShardOptimizer(opt, reduction, name, group_assignment)