提交 fbde872e authored 作者: Brandon T. Willard's avatar Brandon T. Willard 提交者: Brandon T. Willard

Fix ifelse str representation

上级 a1388f92
...@@ -107,7 +107,7 @@ class IfElse(_NoPythonOp): ...@@ -107,7 +107,7 @@ class IfElse(_NoPythonOp):
args.append("inplace") args.append("inplace")
if self.gpu: if self.gpu:
args.append("gpu") args.append("gpu")
return "if{{{','.join(args)}}}" return f"if{{{','.join(args)}}}"
def infer_shape(self, fgraph, node, inputs_shapes): def infer_shape(self, fgraph, node, inputs_shapes):
# By construction, corresponding then/else pairs have the same number # By construction, corresponding then/else pairs have the same number
......
...@@ -520,6 +520,17 @@ class TestIfelse(utt.OptimizationTestMixin): ...@@ -520,6 +520,17 @@ class TestIfelse(utt.OptimizationTestMixin):
loss = ifelse(correct, 0, 1) loss = ifelse(correct, 0, 1)
[(param, param - 0.5 * aesara.grad(cost=loss, wrt=param)) for param in params] [(param, param - 0.5 * aesara.grad(cost=loss, wrt=param)) for param in params]
def test_str(self):
x = vector("x", dtype=self.dtype)
y = vector("y", dtype=self.dtype)
c = iscalar("c")
res = ifelse(c, x, y)
assert str(res.owner).startswith("if{}")
res.owner.op.name = "name"
res.owner.op.as_view = True
res.owner.op.gpu = True
assert str(res.owner).startswith("if{name,inplace,gpu}")
class IfElseIfElseIf(Op): class IfElseIfElseIf(Op):
def __init__(self, inplace=False): def __init__(self, inplace=False):
......
Markdown 格式
0%
您添加了 0 到此讨论。请谨慎行事。
请先完成此评论的编辑!
注册 或者 后发表评论