def sharded_variable(name, shape, num_shards, dtype=tf.float32, transposed=False):
# The final size of the sharded variable may be larger than requested.
# This should be fine for embeddings.
shard_size = int((shape[0] + num_shards - 1) / num_shards)
if transposed:
initializer = tf.uniform_unit_scaling_initializer(dtype=dtype)
else:
initializer = tf.uniform_unit_scaling_initializer(dtype=dtype)
return [tf.get_variable(name + "_" + str(i), [shard_size, shape[1]],
initializer=initializer, dtype=dtype) for i in range(num_shards)]
# XXX(rafal): Code below copied from rnn_cell.py
评论列表
文章目录