import tensorflow as tf

    for k, v in six.iteritems(features):
      v = tf.convert_to_tensor(v)
      v_shape = common_layers.shape_list(v)
      if not v_shape:
        v = tf.expand_dims(v, axis=-1)
        v_shape = [1]
      if v_shape == [1]:
        v = tf.tile(v, [self._num_datashards])
      sharded_features[k] = self._data_parallelism(tf.identity,
                                                   tf.split(
                                                       v, self._num_datashards,
                                                       0))
    return sharded_features

