Set output shape in RANK operation prepare

This fixes segfaults for graphs where RANK output
is a graph internal variable.

Test: GeneratedTest.rank_internal_result
Fix: 150728111
Change-Id: I321d7d4df918d4ecba15af788696f2fbf7545110

1 files changed, 27 insertions, 0 deletions

diff --git a/nn/runtime/test/specs/V1_3/rank.mod.py b/nn/runtime/test/specs/V1_3/rank.mod.py
index e6db5324d..884af35be 100644
--- a/nn/runtime/test/specs/V1_3/rank.mod.py
+++ b/nn/runtime/test/specs/V1_3/rank.mod.py
@@ -41,3 +41,30 @@ test(
     input0_data=[1, 2, 3, 4, 5, 6],
     output0_data=[2],
 )
+
+# b/150728111 regression test.
+# Rank is a first operation that produces a scalar output.
+# This test verifies that RANK works with a scalar output
+# that's internal graph variable (not input or output of a graph).
+def test_internal_output(name, rank_input, fill_dims, fill_output, rank_input_data,
+                  fill_dims_data, fill_output_data):
+  internal_result = Internal("rank_internal_result", "INT32", "{}")
+  model = Model()
+  model = model.Operation("RANK", rank_input).To(internal_result)
+  model = model.Operation("FILL", fill_dims, internal_result).To(fill_output)
+
+  example = Example({
+      rank_input: rank_input_data,
+      fill_dims: fill_dims_data,
+      fill_output: fill_output_data,
+  }, model=model, name=name)
+
+test_internal_output(
+    name="internal_output",
+    rank_input=Input("input0", "TENSOR_FLOAT32", "{2, 3}"),
+    fill_dims=Input("input0", "TENSOR_INT32", "{3}"),
+    fill_output=Output("output", "TENSOR_INT32", "{2, 3, 4}"),
+    rank_input_data=[1, 2, 3, 4, 5, 6],
+    fill_dims_data=[2, 3, 4],
+    fill_output_data=[2] * (2 * 3 * 4),
+)