class_docstr: Partitions a GraphItem's variables according to the given strategy.
  
  Essentially, this does a few things:
  
  1. Reads the Strategy and finds which variables should be partitioned (and how).
  2. Creates new `PartitionedVariables` (essentially a list of `Variables`) and
     replaces the original `Variable` with these new vars. Also splits the original
     gradient to map to each variable shard.
  3. Deletes the original variables and recreates the Optimizer (since Optimizers
     can often have different logic for `PartitionedVariables`, and we don't want
     to rewire the optimizer internals as well).
  4. Returns a new graph and strategy modified to reflect the partitioning.
  
  Note that this currently does not work with variables that are part of a control
  flow (including batchnorm). In those cases, we currently do not have a way to
  replicate the control flow's `FuncGraph` for our new variables. Thus, we expect
  that any StrategyBuilder does not generate a strategy that tries to partition a
  control flow variable, otherwise this will likely error.
class_docstr_tok: ["Partitions","a","GraphItem","'s","variables","according","to","the","given","strategy",".","Essentially",",","this","does","a","few","things",":","1",".","Reads","the","Strategy","and","finds","which","variables","should","be","partitioned","(","and","how",")",".","2",".","Creates","new","`","PartitionedVariables","`","(","essentially","a","list","of","`","Variables","`",")","and","replaces","the","original","`","Variable","`","with","these","new","vars",".","Also","splits","the","original","gradient","to","map","to","each","variable","shard",".","3",".","Deletes","the","original","variables","and","recreates","the","Optimizer","(","since","Optimizers","can","often","have","different","logic","for","`","PartitionedVariables","`",",","and","we","do","n't","want","to","rewire","the","optimizer","internals","as","well",")",".","4",".","Returns","a","new","graph","and","strategy","modified","to","reflect","the","partitioning",".","Note","that","this","currently","does","not","work","with","variables","that","are","part","of","a","control","flow","(","including","batchnorm",")",".","In","those","cases",",","we","currently","do","not","have","a","way","to","replicate","the","control","flow","'s","`","FuncGraph","`","for","our","new","variables",".","Thus",",","we","expect","that","any","StrategyBuilder","does","not","generate","a","strategy","that","tries","to","partition","a","control","flow","variable",",","otherwise","this","will","likely","error","."]
docstr: Constructs new partitioned variables in `new_graph_item`.
  
  Fixes each var's corresponding gradient by splitting the gradient.
  
  Fixes the optimizer by just constructing a new one using the new variables.
  
  Args:
      new_graph_item (GraphItem): The GraphItem in which to construct the new variables and ops.
      vars_to_partition (Dict): Mapping of variable names to number of shards for vars to be partitioned.
      unpartitioned_vars (Dict): Mapping from variable name to gradient name of unpartitioned vars.
  
  Returns:
      List of new variables.
code: def _create_new_vars(
      self, new_graph_item, vars_to_partition, unpartitioned_vars
  ):
      """
      Constructs new partitioned variables in `new_graph_item`.
  
      Fixes each var's corresponding gradient by splitting the gradient.
  
      Fixes the optimizer by just constructing a new one using the new variables.
  
      Args:
          new_graph_item (GraphItem): The GraphItem in which to construct the new variables and ops.
          vars_to_partition (Dict): Mapping of variable names to number of shards for vars to be partitioned.
          unpartitioned_vars (Dict): Mapping from variable name to gradient name of unpartitioned vars.
  
      Returns:
          List of new variables.
      """
      new_grads, new_vars = [], []
      partition_config = {}
      with new_graph_item.graph.as_default():
          for var_name, (
              partition_str,
              reduction_destinations,
          ) in vars_to_partition.items():
              var_op_name = get_op_name(var_name)
              var_op = self.graph_item.graph.get_operation_by_name(
                  var_op_name
              )
              var = self.graph_item.trainable_var_op_to_var[var_op]
              gradient = self.graph_item.var_op_name_to_grad_info[
                  var_op_name
              ][0]
  
              # Create partitioned variable and split gradients
              pc = PartitionerConfig(partition_str=partition_str)
              partition_config[var_op_name] = pc
  
              # Now check compatibility
              if (
                  isinstance(gradient, ops.IndexedSlices)
                  and pc.axis != 0
              ):
                  raise ValueError(
                      "Embedding variables can only be partitioned along the first axis due to "
                      "the limitation on the `embedding_lookup_v2` op."
                  )
  
              initial_value = new_graph_item.graph.get_tensor_by_name(
                  var.initial_value.name
              )
              # NOTE: to enable the saver, for now we only support partition on the one dimension
              # https://github.com/tensorflow/tensorflow/blob/r2.0/tensorflow/python/ops/variables.py#L2915
              partitioned_var = vs.get_variable(
                  var_op.name,
                  shape=None,
                  initializer=initial_value,
                  partitioner=lambda pconf=pc, **unused_kwargs: pconf.partition_list,
                  validate_shape=False,
                  use_resource=True,
              )
              var_list = partitioned_var._variable_list
  
              # Distribute the partitioned variable if they have a PS synchornizer
              # Actually maybe this is not necessary
              for var_slice, device in zip(
                  var_list, reduction_destinations
              ):
                  if device:
                      var_slice.op._set_device_from_string(device)
  
              if isinstance(gradient, ops.IndexedSlices):
                  # Sparse variable
                  new_grad = ops.IndexedSlices(
                      indices=new_graph_item.graph.get_tensor_by_name(
                          gradient.indices.name
                      ),
                      values=new_graph_item.graph.get_tensor_by_name(
                          gradient.values.name
                      ),
                      dense_shape=new_graph_item.graph.get_tensor_by_name(
                          gradient.dense_shape.name
                      ),
                  )
                  split_grad = self._split_indexed_slices_v2(
                      new_grad,
                      len(var_list),
                      var.shape[0],
                      name=f"gradients/splits/sparse_split_{var_op_name}",
                  )
              else:
                  new_grad = new_graph_item.graph.get_tensor_by_name(
                      gradient.name
                  )
  
                  # sometimes new_grad will have polymorphic shape (None), so we use the shape of the original var
                  split_grad = self._split_tensor_v2(
                      new_grad,
                      pc.num_shards,
                      var.shape,
                      pc.axis,
                      name=f"gradients/splits/split_{var_op_name}",
                  )
              self._handle_read(new_graph_item, var_op, partitioned_var)
              self._update_node_config(var, var_list)
  
              self.info.update_variables(var_list, replace=False)
              new_vars.extend(var_list)
              new_grads.extend(split_grad)
              new_graph_item.extend_gradient_info(split_grad, var_list)
              new_graph_item.pop_gradient_info(var.name)
      new_graph_item.info = self.info.copy()
      all_vars, all_grads = new_vars, new_grads
      for var, grad in unpartitioned_vars.items():
          if isinstance(grad, ops.IndexedSlices):
              # Sparse variable
              grad = ops.IndexedSlices(
                  indices=new_graph_item.graph.get_tensor_by_name(
                      grad.indices.name
                  ),
                  values=new_graph_item.graph.get_tensor_by_name(
                      grad.values.name
                  ),
                  dense_shape=new_graph_item.graph.get_tensor_by_name(
                      grad.dense_shape.name
                  ),
              )
          else:
              grad = new_graph_item.graph.get_tensor_by_name(grad.name)
          all_grads.append(grad)
          var = new_graph_item.trainable_var_op_to_var[
              new_graph_item.graph.get_operation_by_name(
                  get_op_name(var)
              )
          ]
          # TensorFlow expects the following to not mess autodist with the tf.distribute
          if (
              not hasattr(var, "_distribute_strategy")
          ) or var._distribute_strategy:
              setattr(var, "_distribute_strategy", None)
          all_vars.append(var)
      with new_graph_item.graph.as_default():
          optimizer = self.graph_item.optimizer(
              *self.graph_item.optimizer_args[1:],
              **self.graph_item.optimizer_kwargs,
          )
          _ = optimizer.apply_gradients(zip(all_grads, all_vars))
      return new_vars, partition_config
  
imports: import collections
  import re
  from google.protobuf.pyext._message import RepeatedScalarContainer
  from tensorflow.core.framework import graph_pb2
  from tensorflow.python import ops
  from tensorflow.python.framework import dtypes, versions
  from tensorflow.python.ops import embedding_ops
  from tensorflow.python.ops import variable_scope, array_ops, math_ops
  from tensorflow.python.ops.variables import Variable, PartitionedVariable
  from tensorflow.python.util.compat import as_bytes
  from autodist.const import AUTODIST_TO_DELETE_SCOPE, COLOCATION_PREFIX
  from autodist.checkpoint.saver import Saver
  from autodist.graph_item import GraphItem, Info
  from autodist.kernel.common.op_info import MUTABLE_STATE_OP_DIRECT_CONSUMER_OPS
  from autodist.kernel.common.utils import get_op_name, get_consumers, update_consumers, parse_name_scope
  from autodist.kernel.common.variable_utils import is_read_var_op
  from autodist.kernel.kernel import Kernel
  from autodist.utils import logging
label_desc: Use image node_id 9 for calling the VariablePartitioner obj's underlying member method code with example usage: obj._create_new_vars(new_graph_item, vars_to_partition, unpartitioned_vars) and returns: new_vars,  partition_config