(self, sum_ops, generators)
| 751 | ONLY_KEEP_IS_AUTO_GEN_SUM_OPS_TAG = "only_keep_is_auto_gen_sum_ops_tag" |
| 752 | |
| 753 | def _SetSumOpsDeviceOption(self, sum_ops, generators): |
| 754 | only_keep_is_auto_gen_sum_ops_tag = False |
| 755 | for generator in generators: |
| 756 | # we already checked that device options are consistent so we can just |
| 757 | # break after finding the first clear_info request |
| 758 | for extra_info in generator.device_option.extra_info: |
| 759 | if extra_info == "{}:1".format(IR.ONLY_KEEP_IS_AUTO_GEN_SUM_OPS_TAG): |
| 760 | only_keep_is_auto_gen_sum_ops_tag = True |
| 761 | break |
| 762 | |
| 763 | if only_keep_is_auto_gen_sum_ops_tag: |
| 764 | # if we find that device_option in the generator that |
| 765 | # requires clear the extra info for the auto gen sum |
| 766 | # Then we will try to clear them and only leave the |
| 767 | # IS_AUTO_GEN_SUM_OPS_TAG |
| 768 | for op in sum_ops: |
| 769 | op.device_option.extra_info.extend([ |
| 770 | "{}:1".format(IR.IS_AUTO_GEN_SUM_OPS_TAG) |
| 771 | ]) |
| 772 | else: |
| 773 | # we already checked that device options are consistent so we can just |
| 774 | # use the first one we find |
| 775 | for generator in generators: |
| 776 | for op in sum_ops: |
| 777 | op.device_option.CopyFrom(generator.device_option) |
| 778 | op.device_option.extra_info.extend([ |
| 779 | "{}:1".format(IR.IS_AUTO_GEN_SUM_OPS_TAG) |
| 780 | ]) |
| 781 | break |
| 782 | |
| 783 | def _DisambiguateGradOpOutput(self, grad_op, idx, cnt): |
| 784 | new_grad_output = ( |
no test coverage detected