class_docstr: Partitioned AR StrategyBuilder.
  
  This StrategyBuilder generates a strategy that partitions each variable along its first dimension,
  and synchronizes them using AllReduce. It might be advantageous for communicating extremely large
  messages -- when synchronizing a single message is bounded by single-flow bandwidth.
  
  It will also sequentially merge collective ops into a single collective group based on chunk_size.
  This strategy does not support synchronizing sparse updates with >1 nodes due to the TF AllGather bug.
docstr: Creates a NodeConfig specifying partitioning and synchronization with AllReduce.
  
  Args:
      var (Variable): The variable to generate a config for.
      var_counter (int): variable counter for collective group assignment.
  
  Returns:
      Dict: the config dict for the node.
code: def _gen_node_config(self, var, var_counter, grad):
      """
      Creates a NodeConfig specifying partitioning and synchronization with AllReduce.
  
      Args:
          var (Variable): The variable to generate a config for.
          var_counter (int): variable counter for collective group assignment.
  
      Returns:
          Dict: the config dict for the node.
      """
      num_shards, partition_axis = self.get_num_shards_and_axis(
          var, grad
      )
      node = strategy_pb2.Strategy.Node()
      node.var_name = var.name
  
      if num_shards <= 1:
          node.AllReduceSynchronizer.spec = (
              synchronizers_pb2.AllReduceSynchronizer.Spec.Value("AUTO")
          )
          node.AllReduceSynchronizer.compressor = (
              synchronizers_pb2.AllReduceSynchronizer.Compressor.Value(
                  "NoneCompressor"
              )
          )
          # node.AllReduceSynchronizer.compressor = \
          #     synchronizers_pb2.AllReduceSynchronizer.Compressor.Value("PowerSGDCompressor")
          node.AllReduceSynchronizer.group = (
              var_counter // self.chunk_size
          )
          return node, num_shards
  
      # num_parts > 1 means the variable will be partitioned
      # generate the partitioner config
      partition_list = [1] * len(var.initial_value.shape)
      partition_list[partition_axis] = num_shards
      pc = PartitionerConfig(partition_list=partition_list)
      node.partitioner = pc.partition_str
      for i in range(num_shards):
          part = strategy_pb2.Strategy.Node()
  
          # If part var_name is inconsistent with what TF will create, partitioner kernel will correct it later.
          # Here let's just make it consistent
          part.var_name = "{}/part_{}:0".format(
              get_op_name(var.name), i
          )
          part.AllReduceSynchronizer.spec = (
              synchronizers_pb2.AllReduceSynchronizer.Spec.Value("AUTO")
          )
          part.AllReduceSynchronizer.compressor = (
              synchronizers_pb2.AllReduceSynchronizer.Compressor.Value(
                  "NoneCompressor"
              )
          )
          # part.AllReduceSynchronizer.compressor = \
          #     synchronizers_pb2.AllReduceSynchronizer.Compressor.Value("PowerSGDCompressor")
          part.AllReduceSynchronizer.group = (
              var_counter + i
          ) // self.chunk_size
          node.part_config.extend([part])
      return node, num_shards
  
file_path: autodist/autodist/strategy/random_axis_partition_all_reduce_strategy.py
imports: import numpy
  from tensorflow.python.framework.ops import IndexedSlices
  from autodist.kernel.common.utils import get_op_name
  from autodist.kernel.partitioner import PartitionerConfig
  from autodist.proto import strategy_pb2, synchronizers_pb2
  from autodist.strategy.base import Strategy, StrategyBuilder
label_desc: Use image node_id 3 for calling the RandomAxisPartitionAR obj's underlying member method code with example usage: obj._gen_node_config(var, var_counter, grad) and returns: node,  num_shards, node,  num_shards