| 789 | REGISTER_OP_GRADIENT("Min", MinGrad); |
| 790 | |
| 791 | static Status MatMulGradHelper(FunctionDef* g, const string& opname, |
| 792 | const string& attr_adj_x, |
| 793 | const string& attr_adj_y, const string& x0, |
| 794 | bool ax0, const string& x1, bool ax1, |
| 795 | const string& y0, bool ay0, const string& y1, |
| 796 | bool ay1, bool enable_broadcasting) { |
| 797 | // The final outputs are "dx" and "dy". If we're broadcasting compute |
| 798 | // intermediate nodes for now. |
| 799 | std::vector<FDH::Node> nodes = { |
| 800 | {{(enable_broadcasting ? "gx" : "dx")}, |
| 801 | opname, |
| 802 | {x0, x1}, |
| 803 | {{"T", "$T"}, {attr_adj_x, ax0}, {attr_adj_y, ax1}}}, |
| 804 | {{(enable_broadcasting ? "gy" : "dy")}, |
| 805 | opname, |
| 806 | {y0, y1}, |
| 807 | {{"T", "$T"}, {attr_adj_x, ay0}, {attr_adj_y, ay1}}}, |
| 808 | }; |
| 809 | // TODO(anudhyan): Figure out a way to inspect the static shapes of "x" and |
| 810 | // "y". If they have the same batch dimensions, then we can omit adding the |
| 811 | // broadcasting-specific ops. |
| 812 | if (enable_broadcasting) { |
| 813 | std::vector<FDH::Node> unbroadcast_gradients = { |
| 814 | FDH::Const<int32>("zero", gtl::ArraySlice<int32>{0}), |
| 815 | FDH::Const<int32>("one", gtl::ArraySlice<int32>{1}), |
| 816 | FDH::Const<int32>("minustwo", gtl::ArraySlice<int32>{-2}), |
| 817 | // Compute the batch shapes of the inputs (all but last two dims). |
| 818 | {{"sx"}, "Shape", {"x"}, {{"T", "$T"}}}, |
| 819 | {{"sy"}, "Shape", {"y"}, {{"T", "$T"}}}, |
| 820 | {{"batch_sx"}, |
| 821 | "StridedSlice", |
| 822 | {"sx", "zero", "minustwo", "one"}, |
| 823 | {{"T", DT_INT32}, {"Index", DT_INT32}}}, |
| 824 | {{"batch_sy"}, |
| 825 | "StridedSlice", |
| 826 | {"sy", "zero", "minustwo", "one"}, |
| 827 | {{"T", DT_INT32}, {"Index", DT_INT32}}}, |
| 828 | // Sum along dimensions that the inputs were broadcasted across. |
| 829 | {{"rx", "ry"}, "BroadcastGradientArgs", {"batch_sx", "batch_sy"}}, |
| 830 | {{"sum_gx"}, "Sum", {"gx", "rx"}, {{"T", "$T"}}}, |
| 831 | {{"sum_gy"}, "Sum", {"gy", "ry"}, {{"T", "$T"}}}, |
| 832 | {{"dx"}, "Reshape", {"sum_gx", "sx"}, {{"T", "$T"}}}, |
| 833 | {{"dy"}, "Reshape", {"sum_gy", "sy"}, {{"T", "$T"}}}}; |
| 834 | nodes.insert(nodes.end(), unbroadcast_gradients.begin(), |
| 835 | unbroadcast_gradients.end()); |
| 836 | } |
| 837 | *g = FDH::Define( |
| 838 | // Arg defs |
| 839 | {"x: T", "y: T", "dz: T"}, |
| 840 | // Ret val defs |
| 841 | {"dx: T", "dy: T"}, |
| 842 | // Attr defs |
| 843 | {{"T: {half, float, double}"}}, |
| 844 | // Nodes |
| 845 | nodes); |
| 846 | return Status::OK(); |
| 847 | } |
| 848 |
no test coverage detected