[GitHub] [tvm] srinidhigoud commented on a change in pull request #7901: [Frontend][Tensorflow] SelectV2 and BroadcastArgs op support for tf2 models

[GitBox <gi...@apache.org>]

srinidhigoud commented on a change in pull request #7901:
URL: https://github.com/apache/tvm/pull/7901#discussion_r617979127



##########
File path: python/tvm/relay/frontend/tensorflow.py
##########
@@ -1765,6 +1765,44 @@ def _impl(inputs, attr, params, mod):
     return _impl
 
 
+def _broadcast_args():
+    def _impl(inputs, attr, params, mod):
+        if isinstance(inputs[0], _expr.Var):
+            s0 = params[inputs[0].name_hint]
+        else:
+            s0 = _infer_value(inputs[0], params, mod)
+        if isinstance(inputs[1], _expr.Var):
+            s1 = params[inputs[1].name_hint]
+        else:
+            s1 = _infer_value(inputs[1], params, mod)
+        s0 = list(s0.asnumpy().reshape([-1]))
+        s1 = list(s1.asnumpy().reshape([-1]))
+        s0_size, s1_size = len(s0), len(s1)
+        from collections import deque
+
+        out = deque([])
+        for i in range(1, min(s0_size, s1_size) + 1):
+            if s0[s0_size - i] == s1[s1_size - i]:
+                out.appendleft(s0[s0_size - i])
+            elif s0[s0_size - i] == 1:
+                out.appendleft(s1[s1_size - i])
+            else:
+                assert s1[s1_size - i] == 1, "Incompatible broadcast type %s and %s" % (
+                    s0[s0_size - i],
+                    s1[s1_size - i],
+                )
+                out.appendleft(s0[s0_size - i])
+        if s0_size < s1_size:
+            for i in range(s0_size + 1, s1_size + 1):
+                out.appendleft(s1[s1_size - i])
+        if s1_size < s0_size:
+            for i in range(s1_size + 1, s0_size + 1):
+                out.appendleft(s0[s0_size - i])

Review comment:
       This will require atleast two list inversions (both the inputs as we have to iterate in reverse to verify broadcasting rules). Fill values will only make sense if I use a zip. It is only two loops (one of last two will fail) and merging them to one will give marginal optimization. 




-- 
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.

For queries about this service, please contact Infrastructure at:
users@infra.apache.org