| 22 | using tensorflow::shape_inference::ShapeHandle; |
| 23 | |
| 24 | Status RpcShapeOp(InferenceContext* c, bool try_rpc) { |
| 25 | ShapeHandle address; |
| 26 | ShapeHandle method; |
| 27 | ShapeHandle request; |
| 28 | ShapeHandle output; |
| 29 | TF_RETURN_IF_ERROR(c->WithRankAtMost(c->input(0), 1, &address)); |
| 30 | if (c->Rank(address) == 1) { |
| 31 | TF_RETURN_IF_ERROR(c->Merge(output, address, &output)); |
| 32 | } |
| 33 | TF_RETURN_IF_ERROR(c->WithRankAtMost(c->input(1), 1, &method)); |
| 34 | if (c->Rank(method) == 1) { |
| 35 | TF_RETURN_IF_ERROR(c->Merge(output, method, &output)); |
| 36 | } |
| 37 | TF_RETURN_IF_ERROR(c->WithRankAtMost(c->input(2), 1, &request)); |
| 38 | if (c->Rank(request) == 1) { |
| 39 | TF_RETURN_IF_ERROR(c->Merge(output, request, &output)); |
| 40 | } |
| 41 | if (!c->RankKnown(output)) { |
| 42 | output = request; |
| 43 | } |
| 44 | c->set_output(0, output); // response |
| 45 | if (try_rpc) { |
| 46 | c->set_output(1, output); // status_code |
| 47 | c->set_output(2, output); // status_message |
| 48 | } |
| 49 | return Status::OK(); |
| 50 | } |
| 51 | |
| 52 | REGISTER_OP("Rpc") |
| 53 | .Input("address: string") |
no test coverage detected