| 513 | }); |
| 514 | |
| 515 | static Status KvResourceApplyAdamShapeFn(InferenceContext* c, bool sparse) { |
| 516 | ShapeHandle unused; |
| 517 | ShapeHandle s = ShapeOrHandleShape(c, 0); // var |
| 518 | TF_RETURN_IF_ERROR(c->Merge(s, ShapeOrHandleShape(c, 1), &s)); // m |
| 519 | TF_RETURN_IF_ERROR(c->Merge(s, ShapeOrHandleShape(c, 2), &s)); // v |
| 520 | TF_RETURN_IF_ERROR(c->WithRank(c->input(3), 0, &unused)); // beta1_power |
| 521 | TF_RETURN_IF_ERROR(c->WithRank(c->input(4), 0, &unused)); // beta2_power |
| 522 | TF_RETURN_IF_ERROR(c->WithRank(c->input(5), 0, &unused)); // lr |
| 523 | TF_RETURN_IF_ERROR(c->WithRank(c->input(6), 0, &unused)); // beta1 |
| 524 | TF_RETURN_IF_ERROR(c->WithRank(c->input(7), 0, &unused)); // beta2 |
| 525 | TF_RETURN_IF_ERROR(c->WithRank(c->input(8), 0, &unused)); // epsilon |
| 526 | TF_RETURN_IF_ERROR( |
| 527 | HandleKvGradAndIndicesInputs(c, sparse, 9 /* grad_idx */, &s)); |
| 528 | if (c->num_outputs() > 0) { |
| 529 | c->set_output(0, s); |
| 530 | } |
| 531 | return Status::OK(); |
| 532 | } |
| 533 | |
| 534 | #define REGISTER_OP_BY_NAME(name) \ |
| 535 | REGISTER_OP(name) \ |
nothing calls this directly
no test coverage detected