提交 af161adb authored 作者: Razvan Pascanu's avatar Razvan Pascanu 提交者: David Warde-Farley

respecting the rval type

As for grad and Rop, I am now taking care that if the input is a list I return a list (regardless of the number of elements).
上级 880db4b2
......@@ -295,6 +295,18 @@ def ifelse(condition, then_branch, else_branch, name=None):
``then_branch`` or in the ``else_branch`` depending on the value of
``cond``.
"""
if type(then_branch) is not type(else_branch):
raise ValueError(('The two branches should be identical. '
'This error could be raised if for example '
' you provided a one element list on the then '
' branch but a tensor on the else branch'))
rval_type = None
if type(then_branch) is list:
rval_type = list
elif type(then_branch) is tuple:
rval_type = tuple
if type(then_branch) not in (list, tuple):
then_branch = [then_branch]
if type(else_branch) not in (list, tuple):
......@@ -314,10 +326,13 @@ def ifelse(condition, then_branch, else_branch, name=None):
ins = [cond] + list(then_branch) + list(else_branch)
rval = new_ifelse.make_node(*ins).outputs
if type(rval) in (list, tuple) and len(rval) == 1:
if rval_type is None:
return rval[0]
elif rval_type is list:
return list(rval)
else:
return rval
return tuple(rval)
@gof.local_optimizer([None])
......
Markdown 格式
0%
您添加了 0 到此讨论。请谨慎行事。
请先完成此评论的编辑!
注册 或者 后发表评论