Get a sharded variable concatenated into one tensor.
(name, shape, dtype, num_shards)
| 48 | |
| 49 | |
| 50 | def _get_concat_variable(name, shape, dtype, num_shards): |
| 51 | """Get a sharded variable concatenated into one tensor.""" |
| 52 | sharded_variable = _get_sharded_variable(name, shape, dtype, num_shards) |
| 53 | if len(sharded_variable) == 1: |
| 54 | return sharded_variable[0] |
| 55 | |
| 56 | concat_name = name + "/concat" |
| 57 | concat_full_name = vs.get_variable_scope().name + "/" + concat_name + ":0" |
| 58 | for value in ops.get_collection(ops.GraphKeys.CONCATENATED_VARIABLES): |
| 59 | if value.name == concat_full_name: |
| 60 | return value |
| 61 | |
| 62 | concat_variable = array_ops.concat(sharded_variable, 0, name=concat_name) |
| 63 | ops.add_to_collection(ops.GraphKeys.CONCATENATED_VARIABLES, concat_variable) |
| 64 | return concat_variable |
| 65 | |
| 66 | |
| 67 | def _get_sharded_variable(name, shape, dtype, num_shards): |
no test coverage detected