class_docstr: An optimizer that averages gradients across TPU shards.
docstr: Verify group_assignment and get the subgroup size".
  
  Args:
    group_assignment: list of group ids for applying the optimizer
      to subgroups.
    num_shards: The number of TPU shards.
  
  Returns:
    The size of one subgroup in group_assignment.
  
  Raises:
    ValueError: If group_assignment is invalid.
docstr_tok: ["Verify","group_assignment","and","get","the","subgroup","size","''",".","Args",":","group_assignment",":","list","of","group","ids","for","applying","the","optimizer","to","subgroups",".","num_shards",":","The","number","of","TPU","shards",".","Returns",":","The","size","of","one","subgroup","in","group_assignment",".","Raises",":","ValueError",":","If","group_assignment","is","invalid","."]
code: def _verify_and_get_subgroup_size(self, group_assignment, num_shards):
      """Verify group_assignment and get the subgroup size".
  
      Args:
        group_assignment: list of group ids for applying the optimizer
          to subgroups.
        num_shards: The number of TPU shards.
  
      Returns:
        The size of one subgroup in group_assignment.
  
      Raises:
        ValueError: If group_assignment is invalid.
      """
      if not group_assignment:
          return None
      if not (
          isinstance(group_assignment, list)
          and all(isinstance(i, list) for i in group_assignment)
      ):
          raise ValueError(
              f"Argument `group_assignment` must be a list of lists. "
              f"Received: {group_assignment}"
          )
  
      replica_ids = set()
      for g in group_assignment:
          for i in g:
              replica_ids.add(i)
  
      if set(range(num_shards)) != replica_ids:
          raise ValueError(
              f"Argument `group_assignment` must be a permutation of "
              f"range({num_shards}). Received: {group_assignment}"
          )
  
      subgroup_size_list = [len(group) for group in group_assignment]
      if all(
          subgroup_size_list[0] == size for size in subgroup_size_list
      ):
          return subgroup_size_list[0]
      else:
          raise ValueError(
              "The size of each subgroup in `group_assignment` must "
              f"be equal. Received: {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 2 for calling the CrossShardOptimizer obj's underlying member method code with example usage: obj._verify_and_get_subgroup_size(group_assignment, num_shards) and returns: None, subgroup_size_list