Get a list of sharded variables with the given dtype.
(name, shape, dtype, num_shards)
| 65 | |
| 66 | |
| 67 | def _get_sharded_variable(name, shape, dtype, num_shards): |
| 68 | """Get a list of sharded variables with the given dtype.""" |
| 69 | if num_shards > shape[0]: |
| 70 | raise ValueError("Too many shards: shape=%s, num_shards=%d" % (shape, |
| 71 | num_shards)) |
| 72 | unit_shard_size = int(math.floor(shape[0] / num_shards)) |
| 73 | remaining_rows = shape[0] - unit_shard_size * num_shards |
| 74 | |
| 75 | shards = [] |
| 76 | for i in range(num_shards): |
| 77 | current_size = unit_shard_size |
| 78 | if i < remaining_rows: |
| 79 | current_size += 1 |
| 80 | shards.append( |
| 81 | vs.get_variable( |
| 82 | name + "_%d" % i, [current_size] + shape[1:], dtype=dtype)) |
| 83 | return shards |
| 84 | |
| 85 | |
| 86 | def _norm(g, b, inp, scope): |
no test coverage detected