Compute the output shape of the merge node as the union of the available input shapes.
| 1878 | // Compute the output shape of the merge node as the union of the available |
| 1879 | // input shapes. |
| 1880 | Status GraphProperties::UpdateMerge(SymbolicShapeRefiner* shape_refiner, |
| 1881 | const NodeDef* node, |
| 1882 | bool* new_shapes) const { |
| 1883 | InferenceContext* ic = shape_refiner->GetContext(node); |
| 1884 | if (!ic) { |
| 1885 | // Now we can run shape inference |
| 1886 | TF_RETURN_IF_ERROR(shape_refiner->AddNode(node)); |
| 1887 | ic = CHECK_NOTNULL(shape_refiner->GetContext(node)); |
| 1888 | *new_shapes = true; |
| 1889 | |
| 1890 | // Infer the shape of the second output once and for all since it never |
| 1891 | // changes. |
| 1892 | ShapeHandle out1 = ic->Scalar(); |
| 1893 | ic->set_output(1, out1); |
| 1894 | } |
| 1895 | |
| 1896 | ShapeHandle out; |
| 1897 | const std::vector<ShapeAndType>* out_handle = nullptr; |
| 1898 | bool out_initialized = false; |
| 1899 | for (const GraphView::Edge fanin : shape_refiner->graph().GetFaninEdges( |
| 1900 | *node, /*include_controlling_edges=*/false)) { |
| 1901 | InferenceContext* src_ic = shape_refiner->GetContext(fanin.src.node); |
| 1902 | if (!src_ic) { |
| 1903 | // Handling a loop for the first time, the back edge won't have any shape |
| 1904 | // info. |
| 1905 | continue; |
| 1906 | } |
| 1907 | ShapeHandle input = src_ic->output(fanin.src.port_id); |
| 1908 | ic->SetInput(fanin.dst.port_id, input); |
| 1909 | auto* input_handle = |
| 1910 | src_ic->output_handle_shapes_and_types(fanin.src.port_id); |
| 1911 | if (input_handle) |
| 1912 | ic->set_input_handle_shapes_and_types(fanin.dst.port_id, *input_handle); |
| 1913 | if (!out_initialized) { |
| 1914 | out_initialized = true; |
| 1915 | out = input; |
| 1916 | out_handle = input_handle; |
| 1917 | } else { |
| 1918 | // Note here only out, not out_handle, is modified. |
| 1919 | out = shape_refiner->OutputAsUnion(node, 0, input, out); |
| 1920 | } |
| 1921 | } |
| 1922 | |
| 1923 | if (*new_shapes || !shape_refiner->EquivalentShapes(out, ic->output(0))) { |
| 1924 | ic->set_output(0, out); |
| 1925 | if (out_handle) ic->set_output_handle_shapes_and_types(0, *out_handle); |
| 1926 | *new_shapes = true; |
| 1927 | } |
| 1928 | |
| 1929 | return Status::OK(); |
| 1930 | } |
| 1931 | |
| 1932 | // Manually propagate the input shape for Enter nodes. |
| 1933 | Status GraphProperties::UpdateEnter(SymbolicShapeRefiner* shape_refiner, |
nothing calls this directly
no test coverage detected