This optimizer first rewrites Prod(Shape(x)) into Size(x). It then uses symbolic shapes to simplify Div(Size(x), Size(y)) in the case that x and y share symbolic shapes that are unknown but known to be identical, e.g. we can deduce that Div(Size([2,?,2]) Size([1,?,2])) is 2 if the two unknown dimensions are known to be identical. This can be inferred if they share the same symbolic representation
| 35 | // dimensions are known to be identical. This can be inferred if they share the |
| 36 | // same symbolic representation (negative integer dimension size). |
| 37 | Status ShapeOptimizer::Optimize(Cluster* cluster, const GrapplerItem& item, |
| 38 | GraphDef* optimized_graph) { |
| 39 | // Do a quick check to determine if we can skip this optimizer. |
| 40 | bool can_optimize = false; |
| 41 | bool has_div = false; |
| 42 | bool has_size = false; |
| 43 | bool has_shape = false; |
| 44 | bool has_prod = false; |
| 45 | auto is_int = [](const NodeDef& node) -> bool { |
| 46 | return node.attr().at("T").type() == DT_INT32 || |
| 47 | node.attr().at("T").type() == DT_INT64; |
| 48 | }; |
| 49 | for (const NodeDef& node : item.graph.node()) { |
| 50 | if (IsShape(node)) { |
| 51 | has_shape = true; |
| 52 | } else if (IsProd(node) && is_int(node)) { |
| 53 | has_prod = true; |
| 54 | } else if (IsDiv(node) && is_int(node)) { |
| 55 | has_div = true; |
| 56 | } else if (IsSize(node)) { |
| 57 | has_size = true; |
| 58 | } |
| 59 | if ((has_shape && has_prod) || (has_div && has_size)) { |
| 60 | can_optimize = true; |
| 61 | break; |
| 62 | } |
| 63 | } |
| 64 | if (!can_optimize) { |
| 65 | return errors::Aborted("Nothing to do."); |
| 66 | } |
| 67 | |
| 68 | *optimized_graph = item.graph; |
| 69 | MutableGraphView graph(optimized_graph); |
| 70 | GraphProperties properties(item); |
| 71 | bool inferred_properties = false; |
| 72 | |
| 73 | // The product of all the dimensions in a tensor shape can be expressed more |
| 74 | // simply as the size of the tensor. |
| 75 | for (auto& node : *optimized_graph->mutable_node()) { |
| 76 | if (!IsShape(node)) { |
| 77 | continue; |
| 78 | } |
| 79 | for (MutableGraphView::InputPort fanout : |
| 80 | graph.GetFanout(MutableGraphView::OutputPort(&node, 0))) { |
| 81 | if (fanout.node->op() != "Prod") { |
| 82 | continue; |
| 83 | } |
| 84 | if (fanout.node->attr().count("keep_dims") != 0 && |
| 85 | fanout.node->attr().at("keep_dims").b()) { |
| 86 | // Keeping the reduced dimensions won't result in a scalar, so we can't |
| 87 | // rewrite the whole expression directly as a Size operation. |
| 88 | continue; |
| 89 | } |
| 90 | const MutableGraphView::OutputPort reduce_indices = |
| 91 | graph.GetRegularFanin(MutableGraphView::InputPort(fanout.node, 1)); |
| 92 | if (!inferred_properties) { |
| 93 | // Infer properties lazily in case they are not needed. |
| 94 | TF_RETURN_IF_ERROR( |
nothing calls this directly
no test coverage detected