(self, sparse_generators)
| 625 | ) |
| 626 | |
| 627 | def AppendSparseGenerators(self, sparse_generators): |
| 628 | # merge indices and values generators for sparse gradients |
| 629 | for name, input_generators in sparse_generators.items(): |
| 630 | for version, generators in input_generators.items(): |
| 631 | if len(generators) == 1: |
| 632 | # either indices or values are generated (but not both) |
| 633 | generator = generators[0] |
| 634 | else: |
| 635 | # both indices and values are generated |
| 636 | assert(len(generators) == 2) |
| 637 | op1_i, idx1_i, op1_v, idx1_v, g1, dev_1 = generators[0] |
| 638 | op2_i, idx2_i, op2_v, idx2_v, g2, dev_2 = generators[1] |
| 639 | assert(g1 == g2) |
| 640 | assert dev_1 == dev_2, ( |
| 641 | "Unequal devices for sparse generators: " |
| 642 | "{} and {}".format(dev_1, dev_2) |
| 643 | ) |
| 644 | assert(op1_i is None or op2_i is None) |
| 645 | assert(op1_v is None or op2_v is None) |
| 646 | assert(idx1_i == 0 or idx2_i == 0) |
| 647 | assert(idx1_v == 0 or idx2_v == 0) |
| 648 | generator = SparseGradGenMeta( |
| 649 | op1_i or op2_i, idx1_i + idx2_i, |
| 650 | op1_v or op2_v, idx1_v + idx2_v, |
| 651 | g1, dev_1) |
| 652 | self.gradient_generators[name][version].append(generator) |
| 653 | |
| 654 | def BuildGradientGenerators( # NOQA |
| 655 | self, fwd_op_idx, gradient_ops, g_output, g_input): |
no test coverage detected