Improve static shape inference in grappler by propagating tensors_as_shapes better:
Currently, static shape inference propagates shapes of tensors, but in some cases, we do need values; for this, we use input_tensors (from Const input tensor) and input_tensors_as_shapes and output_tensors_as_shapes (these are ShapeHandle format, but has values, currently only for 1D vector). This CL enhances propagation of input_tensors_as_shapes and output_tensors_as_shapes to improve static shape inference. (1) forward scalar Const as input_tensors_as_shapes (currently, only 1D vector), (2) export input_tensors_as_shapes, output const tensor, and output_tensors_as_shapes to the values of inferred input/output TensorProperties (currently, only input const tensors are exported as values), (3) use input_tensors_as_shapes as Const tensor to function input (currently, only Const tensors), (4) forward input_tensors_as_shapes to output_tensors_as_shapes for Identity op, (5) when Pack op concats scalar values to form output_tensors_as_shapes, currently it uses only input_tensors (from Const input tensors), but this CL change Pack to use input_tensors_as_shapes as well. PiperOrigin-RevId: 212696959
Loading
Please sign in to comment