From 67a6f7f26f092860a4e1b3f6ac0c2e5cde9bf685 Mon Sep 17 00:00:00 2001
From: Dominic Symes <dominic.symes@arm.com>
Date: Mon, 6 Mar 2023 17:14:15 +0000
Subject: Check the output shape of broadcast operations

For an operation that performs broadcast the
output shape size must be the maximum of the input shape
sizes in each dimension.

Additionally, the input dimension size must be 1 whenever
an input shape does not match the output shape size
in a dimension.

Signed-off-by: Dominic Symes <dominic.symes@arm.com>
Change-Id: I89492f4ef22da76f84f12e720c79634ea42545bc
---
 chapters/pseudocode.adoc | 20 +++++++++++++++++++-
 1 file changed, 19 insertions(+), 1 deletion(-)

(limited to 'chapters/pseudocode.adoc')

diff --git a/chapters/pseudocode.adoc b/chapters/pseudocode.adoc
index db699d1..42f123b 100644
--- a/chapters/pseudocode.adoc
+++ b/chapters/pseudocode.adoc
@@ -125,7 +125,25 @@ void tensor_write<type>(<type> *address, dim_t shape, dim_t index, <type> value)
 }
 ----
 
-==== Broadcast Helper
+==== Broadcast Helpers
+
+The following function derives the broadcast output shape from the input shapes.
+
+[source,c++]
+----
+dim_t broadcast_shape(dim_t shape1, dim_t shape2) {
+    ERROR_IF(rank(shape1) != rank(shape2));
+    dim_t shape = shape1;
+    for (int32_t i = 0; i < rank(shape); i++) {
+        if (shape[i] == 1) {
+            shape[i] = shape2[i];
+        } else {
+            ERROR_IF(shape2[i] != 1 && shape2[i] != shape[i]);
+        }
+    }
+    return shape;
+}
+----
 
 The following function maps an index in the output tensor to an index in the input tensor.
 
-- 
cgit v1.2.1

