| 91 | )doc"); |
| 92 | |
| 93 | Status IfShapeInferenceFn(shape_inference::InferenceContext* c) { |
| 94 | std::vector<PartialTensorShape> output_shapes; |
| 95 | TF_RETURN_IF_ERROR(c->GetAttr("output_shapes", &output_shapes)); |
| 96 | // If `output_shapes` attr is set use that as the shapes of the outputs |
| 97 | // else return unknown shapes. |
| 98 | if (output_shapes.empty()) return shape_inference::UnknownShape(c); |
| 99 | if (output_shapes.size() != c->num_outputs()) { |
| 100 | return errors::InvalidArgument( |
| 101 | "`output_shapes` must be the same length as num outputs (", |
| 102 | output_shapes.size(), " vs. ", c->num_outputs()); |
| 103 | } |
| 104 | for (size_t i = 0; i < output_shapes.size(); ++i) { |
| 105 | shape_inference::ShapeHandle output_shape_handle; |
| 106 | TF_RETURN_IF_ERROR(c->MakeShapeFromPartialTensorShape( |
| 107 | output_shapes[i], &output_shape_handle)); |
| 108 | c->set_output(static_cast<int>(i), output_shape_handle); |
| 109 | } |
| 110 | return Status::OK(); |
| 111 | } |
| 112 | |
| 113 | REGISTER_OP("StatelessIf") |
| 114 | .Input("cond: Tcond") |
nothing calls this directly
no test coverage detected