| 22 | |
| 23 | |
| 24 | def embedding_update(params, |
| 25 | ids, |
| 26 | values, |
| 27 | partition_strategy='mod', |
| 28 | func=tf.scatter_update, |
| 29 | name=None): |
| 30 | if isinstance(params, variables.PartitionedVariable): |
| 31 | params = list(params) |
| 32 | if not isinstance(params, list): |
| 33 | params = [params] |
| 34 | |
| 35 | params = list(params) |
| 36 | np = len(params) |
| 37 | |
| 38 | if np == 1: |
| 39 | with tf.colocate_with(params[0]): |
| 40 | return func(params[0], ids, values) |
| 41 | |
| 42 | if partition_strategy == 'mod': |
| 43 | p_assignments = ids % np |
| 44 | new_ids = ids // np |
| 45 | else: |
| 46 | raise ValueError('Unrecognized partition strategy: ' + |
| 47 | partition_strategy) |
| 48 | |
| 49 | p_assignments = tf.to_int32(p_assignments) |
| 50 | scatter_ids = tf.dynamic_partition(new_ids, p_assignments, np) |
| 51 | scatter_values = tf.dynamic_partition(values, p_assignments, np) |
| 52 | |
| 53 | update_ops = [] |
| 54 | for param, pids, pvalues in zip(params, scatter_ids, scatter_values): |
| 55 | with tf.colocate_with(param): |
| 56 | update_ops.append(func(param, pids, pvalues)) |
| 57 | |
| 58 | return tf.group(*update_ops) |
| 59 | |
| 60 | |
| 61 | def embedding_add(params, |