| 587 | #undef REGISTER_OP_BY_NAME |
| 588 | |
| 589 | static Status KvApplyAdamAsyncShapeFn(InferenceContext* c, bool sparse) { |
| 590 | ShapeHandle unused; |
| 591 | ShapeHandle s = ShapeOrHandleShape(c, 0); // var |
| 592 | TF_RETURN_IF_ERROR(c->Merge(s, ShapeOrHandleShape(c, 1), &s)); // m |
| 593 | TF_RETURN_IF_ERROR(c->Merge(s, ShapeOrHandleShape(c, 2), &s)); // v |
| 594 | TF_RETURN_IF_ERROR(c->WithRank(c->input(3), 0, &unused)); // beta1_power |
| 595 | TF_RETURN_IF_ERROR(c->WithRank(c->input(4), 0, &unused)); // beta2_power |
| 596 | TF_RETURN_IF_ERROR(c->WithRank(c->input(5), 0, &unused)); // lr |
| 597 | TF_RETURN_IF_ERROR(c->WithRank(c->input(6), 0, &unused)); // beta1 |
| 598 | TF_RETURN_IF_ERROR(c->WithRank(c->input(7), 0, &unused)); // beta2 |
| 599 | TF_RETURN_IF_ERROR(c->WithRank(c->input(8), 0, &unused)); // epsilon |
| 600 | TF_RETURN_IF_ERROR( |
| 601 | HandleKvGradAndIndicesInputs(c, sparse, 9 /* grad_idx */, &s)); |
| 602 | if (c->num_outputs() > 0) { |
| 603 | c->set_output(0, s); |
| 604 | } |
| 605 | return Status::OK(); |
| 606 | } |
| 607 | |
| 608 | #define REGISTER_OP_BY_NAME(name) \ |
| 609 | REGISTER_OP(name) \ |
nothing calls this directly
no test coverage detected