Updates t-th row in accumulators. Args: struct_acc: The accumulators. A structure of tensors. struct_x: The new values. A structure of tensors congruent to `struct_acc`. t: A scalar integer. Performance is better if `t` is on the device memory. Returns: A structure of ten
(struct_acc, struct_x, t)
| 83 | |
| 84 | |
| 85 | def _Update(struct_acc, struct_x, t): |
| 86 | """Updates t-th row in accumulators. |
| 87 | |
| 88 | Args: |
| 89 | struct_acc: The accumulators. A structure of tensors. |
| 90 | struct_x: The new values. A structure of tensors congruent to `struct_acc`. |
| 91 | t: A scalar integer. Performance is better if `t` is on the device |
| 92 | memory. |
| 93 | |
| 94 | Returns: |
| 95 | A structure of tensors. Say, ret is a returned dictionary. Then, for |
| 96 | each key, we have: |
| 97 | ret[key] = struct_acc[key]; |
| 98 | ret[key][t, :] = struct_x[key] |
| 99 | """ |
| 100 | to_skip_update = set() |
| 101 | acc_lst = nest.flatten(struct_acc) |
| 102 | x_lst = nest.flatten(struct_x) |
| 103 | t = math_ops.cast( |
| 104 | [t], dtypes.int32) # tf.compat.v1.to_int32 casts on-device tensors. |
| 105 | lst = [] |
| 106 | for acc, x in zip(acc_lst, x_lst): |
| 107 | if acc in to_skip_update: |
| 108 | # Until b/62105730 is fixed, we need to avoid inplace update for tensors |
| 109 | # of rank 1. could reshape to handle it, but we don't really need the |
| 110 | # values applied to these, so just skip their modification. |
| 111 | lst += [acc] |
| 112 | else: |
| 113 | lst += [alias_inplace_update(acc, t, array_ops.expand_dims(x, 0))] |
| 114 | return nest.pack_sequence_as(struct_acc, lst) |
| 115 | |
| 116 | |
| 117 | def _SeqLenDim(struct): |
no test coverage detected