rohanmukh commented on a change in pull request #7901:
URL: https://github.com/apache/tvm/pull/7901#discussion_r617970021
##########
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:
Is it better to use itertools.zip_longest() on the reversed array
instead of running three separate loops? Might be helpful to have fill_values
set at 1 to avoid handling None values.
--
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:
[email protected]
