def create_convert_ifelse_node(return_name_ids,
pred,
true_func,
false_func,
is_if_expr=False):
"""
Create `paddle.jit.dy2static.convert_ifelse(
pred, true_fn, false_fn, true_args, false_args, return_vars)`
to replace original `python if/else` statement.
"""
def create_name_nodes(name_ids):
if not name_ids:
return gast.Tuple(elts=[], ctx=gast.Load())
gast_names = [
gast.Name(id=name_id,
ctx=gast.Load(),
annotation=None,
type_comment=None) for name_id in name_ids
]
name_node = gast.Tuple(elts=gast_names, ctx=gast.Load())
return name_node
if is_if_expr:
true_args = gast.Tuple(elts=[], ctx=gast.Load())
false_args = gast.Tuple(elts=[], ctx=gast.Load())
true_func_source = "lambda : {}".format(ast_to_source_code(true_func))
false_func_source = "lambda : {}".format(
ast_to_source_code(false_func))
else:
true_args = gast.Tuple(elts=true_func.args.args, ctx=gast.Load())
false_args = gast.Tuple(elts=false_func.args.args, ctx=gast.Load())
true_func_source = true_func.name
false_func_source = false_func.name
return_vars = create_name_nodes(return_name_ids)
convert_ifelse_layer = gast.parse(
'_jst.convert_ifelse('
'{pred}, {true_fn}, {false_fn}, {true_args}, {false_args}, {return_vars})'
.format(pred=ast_to_source_code(pred),
true_fn=true_func_source,
false_fn=false_func_source,
true_args=ast_to_source_code(true_args),
false_args=ast_to_source_code(false_args),
return_vars=ast_to_source_code(return_vars))).body[0].value
if return_name_ids:
_, cond_node = create_assign_node(return_name_ids,
convert_ifelse_layer)
else: # No variables can be returned if no assign statement in if.body.
cond_node = gast.Expr(value=convert_ifelse_layer)
return cond_node
def _create_node_with_api_template(self, node, template):
node_code = ast_to_source_code(node)
new_node_str = template.format(node_code)
# gast.parse return Module(body=[expr(value=...)])
new_node = gast.parse(new_node_str).body[0].value
bool_tensor_name = unique_name.generate(PLAIN_TENSOR_PREFIX)
assign_name, assign_node = create_assign_node(bool_tensor_name,
new_node)
self._new_assign_nodes.append(assign_node)
return assign_name
def _create_bool_node(self, node):
node_code = ast_to_source_code(node)
new_node_str = "fluid.layers.fill_constant(shape=[1], dtype='bool', value=bool({}))".format(
node_code)
# gast.parse return Module(body=[expr(value=...)])
new_node = gast.parse(new_node_str).body[0].value
bool_tensor_name = unique_name.generate(PLAIN_TENSOR_PREFIX)
assign_name, assign_node = create_assign_node(bool_tensor_name,
new_node)
self._new_assign_nodes.append(assign_node)
return assign_name
def visit_UnaryOp(self, node):
self.generic_visit(node)
if isinstance(node.op, gast.Not):
arg = ast_to_source_code(node.operand)
new_node_str = "fluid.layers.logical_not({})".format(arg)
# gast.parse returns Module(body=[expr(value=...)])
new_node = gast.parse(new_node_str).body[0].value
logic_tensor_name = unique_name.generate(LOGIC_NOT_PREFIX)
assign_name, assign_node = create_assign_node(
logic_tensor_name, new_node)
self._new_assign_nodes.append(assign_node)
return assign_name
return node
def create_cond_node(return_name_ids,
pred,
true_func,
false_func,
is_if_expr=False):
"""
Create `fluid.layers.cond(pred, true_fn, false_fn)` to replace
original `python if/else` statement.
"""
def create_lambda_node(func_or_expr_node, is_if_expr=False):
body = func_or_expr_node
if not is_if_expr:
body = gast.Call(func=gast.Name(id=func_or_expr_node.name,
ctx=gast.Load(),
annotation=None,
type_comment=None),
args=[func_or_expr_node.args],
keywords=[])
lambda_node = gast.Lambda(args=gast.arguments(args=[],
posonlyargs=[],
vararg=None,
kwonlyargs=[],
kw_defaults=None,
kwarg=None,
defaults=[]),
body=body)
return lambda_node
cond_api = gast.parse('fluid.layers.cond').body[0].value
true_func_lambda = create_lambda_node(true_func, is_if_expr)
false_func_lambda = create_lambda_node(false_func, is_if_expr)
cond_layer = gast.Call(func=cond_api,
args=[pred, true_func_lambda, false_func_lambda],
keywords=[])
if return_name_ids:
_, cond_node = create_assign_node(return_name_ids, cond_layer)
else: # No variables can be returned if no assign statement in if.body.
cond_node = gast.Expr(value=cond_layer)
return cond_node
def _create_node(nodes, api_type):
assert len(
nodes
) > 1, "The length of BoolOp should be at least 2, but received {}.".format(
len(nodes))
if len(nodes) > 2:
# Creates logic_and/logic_or node recursively.
pre_assign_node = _create_node(nodes[:2], api_type)
nodes = [pre_assign_node] + nodes[2:]
args = [ast_to_source_code(child) for child in nodes]
new_node_str = "fluid.layers.logical_{}(x={}, y={})".format(
api_type, args[0], args[1])
# gast.parse return Module(body=[expr(value=...)])
new_node = gast.parse(new_node_str).body[0].value
logic_tensor_name = unique_name.generate(
LOGIC_AND_PREFIX if 'and' in api_type else LOGIC_OR_PREFIX)
assign_name, assign_node = create_assign_node(
logic_tensor_name, new_node)
self._new_assign_nodes.append(assign_node)
return assign_name