from tensorflow.contrib import layers

  parent_scope = "dnn"

  input_layer_partitioner = (partitioned_variables.min_max_variable_partitioner(
      max_partitions=num_ps_replicas, min_slice_size=64 << 20))
  input_layer_scope = parent_scope + "/input_from_feature_columns"
  with variable_scope.variable_scope(
      input_layer_scope,
      values=list(six.itervalues(features)),
      partitioner=input_layer_partitioner) as scope:
    net = layers.input_from_feature_columns(
        columns_to_tensors=features,
        feature_columns=feature_columns,
        weight_collections=[parent_scope],
        scope=scope)

  hidden_layer_partitioner = (
      partitioned_variables.min_max_variable_partitioner(
          max_partitions=num_ps_replicas))
