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

2 files changed, 98 insertions, 0 deletions

diff --git a/nn/runtime/test/generated/spec_V1_3/rank.example.cpp b/nn/runtime/test/generated/spec_V1_3/rank.example.cpp
index 69db68b93..6e0fa86b1 100644
--- a/nn/runtime/test/generated/spec_V1_3/rank.example.cpp
+++ b/nn/runtime/test/generated/spec_V1_3/rank.example.cpp
@@ -1122,3 +1122,74 @@ const auto dummy_test_model_1d_quant8_signed_all_inputs_as_internal_2 = TestMode
 
 }  // namespace generated_tests::rank
 
+namespace generated_tests::rank {
+
+const TestModel& get_test_model_internal_output() {
+    static TestModel model = {
+        .expectFailure = false,
+        .expectedMultinomialDistributionTolerance = 0,
+        .isRelaxed = false,
+        .main = {
+                .inputIndexes = {0, 2},
+                .operands = {{
+                            .channelQuant = {},
+                            .data = TestBuffer::createFromVector<float>({1.0f, 2.0f, 3.0f, 4.0f, 5.0f, 6.0f}),
+                            .dimensions = {2, 3},
+                            .isIgnored = false,
+                            .lifetime = TestOperandLifeTime::SUBGRAPH_INPUT,
+                            .numberOfConsumers = 1,
+                            .scale = 0.0f,
+                            .type = TestOperandType::TENSOR_FLOAT32,
+                            .zeroPoint = 0
+                        }, {
+                            .channelQuant = {},
+                            .data = TestBuffer::createFromVector<int32_t>({}),
+                            .dimensions = {},
+                            .isIgnored = false,
+                            .lifetime = TestOperandLifeTime::TEMPORARY_VARIABLE,
+                            .numberOfConsumers = 1,
+                            .scale = 0.0f,
+                            .type = TestOperandType::INT32,
+                            .zeroPoint = 0
+                        }, {
+                            .channelQuant = {},
+                            .data = TestBuffer::createFromVector<int32_t>({2, 3, 4}),
+                            .dimensions = {3},
+                            .isIgnored = false,
+                            .lifetime = TestOperandLifeTime::SUBGRAPH_INPUT,
+                            .numberOfConsumers = 1,
+                            .scale = 0.0f,
+                            .type = TestOperandType::TENSOR_INT32,
+                            .zeroPoint = 0
+                        }, {
+                            .channelQuant = {},
+                            .data = TestBuffer::createFromVector<int32_t>({2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2}),
+                            .dimensions = {2, 3, 4},
+                            .isIgnored = false,
+                            .lifetime = TestOperandLifeTime::SUBGRAPH_OUTPUT,
+                            .numberOfConsumers = 0,
+                            .scale = 0.0f,
+                            .type = TestOperandType::TENSOR_INT32,
+                            .zeroPoint = 0
+                        }},
+                .operations = {{
+                            .inputs = {0},
+                            .outputs = {1},
+                            .type = TestOperationType::RANK
+                        }, {
+                            .inputs = {2, 1},
+                            .outputs = {3},
+                            .type = TestOperationType::FILL
+                        }},
+                .outputIndexes = {3}
+            },
+        .minSupportedVersion = TestHalVersion::V1_3,
+        .referenced = {}
+    };
+    return model;
+}
+
+const auto dummy_test_model_internal_output = TestModelManager::get().add("rank_internal_output", get_test_model_internal_output());
+
+}  // namespace generated_tests::rank
+
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),
+)