message: Disable variable partitioning for moving_mean and moving_variance for fused batch norm.
diff: @@ -350,7 +350,13 @@ def _fused_batch_norm(
  gamma = array_ops.constant(1.0, shape=params_shape)
  # Create moving_mean and moving_variance variables and add them to the
  - # appropriate collections.
  + # appropriate collections. We disable variable partitioning while creating
  + # them, because assign_moving_average is not yet supported for partitioned
  + # variables (this needs to be handled carefully, as it may break
  + # the checkpoint backward compatibility).
  + with variable_scope.variable_scope(
  + variable_scope.get_variable_scope()) as local_scope:
  + local_scope.set_partitioner(None)
  moving_mean_collections = utils.get_variable_collections(
  variables_collections, 'moving_mean')
  moving_mean_initializer = param_initializers.get(
  @@ -717,10 +723,11 @@ def batch_norm(inputs,
  # Create moving_mean and moving_variance variables and add them to the
  # appropriate collections. We disable variable partitioning while creating
  # them, because assign_moving_average is not yet supported for partitioned
  - # variables.
  - partitioner = variable_scope.get_variable_scope().partitioner
  - try:
  - variable_scope.get_variable_scope().set_partitioner(None)
  + # variables (this needs to be handled carefully, as it may break
  + # the checkpoint backward compatibility).
  + with variable_scope.variable_scope(
  + variable_scope.get_variable_scope()) as local_scope:
  + local_scope.set_partitioner(None)
  moving_mean_collections = utils.get_variable_collections(
  variables_collections, 'moving_mean')
  moving_mean_initializer = param_initializers.get(
  @@ -743,8 +750,6 @@ def batch_norm(inputs,
  initializer=moving_variance_initializer,
  trainable=False,
  collections=moving_variance_collections)
  - finally:
  - variable_scope.get_variable_scope().set_partitioner(partitioner)
  # If `is_training` doesn't have a constant value, because it is a `Tensor`,
  # a `Variable` or `Placeholder` then is_training_value will be None and
  