TODO(agarwal): use an automatic mechanism for handling None arguments to gradient functions. Some gradient functions can accept None arguments for gradients. The following maps the operation name to the indices at which the corresponding gradient function can accept None values. e.g. FusedBatchNorm outputs 5 values and hence receives 5 gradient values during backprop. However the gradient functio
| 609 | // ignored (and hence can be None). The backprop algorithm can then leverage |
| 610 | // this by not constructing zeros to pass for those indices. |
| 611 | std::unordered_map<string, std::unordered_set<int>>* |
| 612 | FunctionsAcceptingNoneForIndicesMap() { |
| 613 | static auto* const m = |
| 614 | new std::unordered_map<string, std::unordered_set<int>>({ |
| 615 | {"SoftmaxCrossEntropyWithLogits", {1}}, |
| 616 | {"SparseSoftmaxCrossEntropyWithLogits", {1}}, |
| 617 | {"FusedBatchNorm", {1, 2, 3, 4}}, |
| 618 | }); |
| 619 | return m; |
| 620 | } |
| 621 | |
| 622 | } // namespace |
| 623 |