Cwise ternary ops.
| 645 | |
| 646 | // Cwise ternary ops. |
| 647 | Status SelectGrad(const AttrSlice& attrs, FunctionDef* g) { |
| 648 | // clang-format off |
| 649 | *g = FDH::Define( |
| 650 | {"c:bool", "x:T", "y:T", "dz:T"}, |
| 651 | {"dc:bool", "dx:T", "dy:T"}, |
| 652 | {{"T: {half, float, double}"}}, |
| 653 | { |
| 654 | {{"dc"}, "ZerosLike", {"c"}, {{"T", DT_BOOL}}, {"dz"}}, |
| 655 | {{"zeros"}, "ZerosLike", {"x"}, {{"T", "$T"}}, {"dz"}}, |
| 656 | {{"dx"}, "Select", {"c", "dz", "zeros"}, {{"T", "$T"}}}, |
| 657 | {{"dy"}, "Select", {"c", "zeros", "dz"}, {{"T", "$T"}}}, |
| 658 | }); |
| 659 | // clang-format on |
| 660 | return Status::OK(); |
| 661 | } |
| 662 | REGISTER_OP_GRADIENT("Select", SelectGrad); |
| 663 | |
| 664 | // N-ry ops |