def _update_name_to_var_shape(self, node):
assert isinstance(node, gast.Assign)
target_node = node.targets[0]
value_node = node.value
if isinstance(target_node, gast.Tuple):
has_updated = False
for idx, element in enumerate(target_node.elts):
target_id = ast_to_source_code(element).strip()
if isinstance(value_node, gast.Name):
if value_node.id in self.name_to_var_shape:
index_value_node = gast.Constant(value=idx, kind=None)
slice_index_node = gast.Index(value=index_value_node)
var_shape_node = self.name_to_var_shape[value_node.id]
sub_node = gast.Subscript(value=var_shape_node,
slice=slice_index_node,
ctx=gast.Load())
self.name_to_var_shape[target_id] = sub_node
has_updated = True
if isinstance(value_node, gast.Attribute):
if self.is_var_shape(value_node): # eg: x.shape
index_value_node = gast.Constant(value=idx, kind=None)
slice_index_node = gast.Index(value=index_value_node)
sub_node = gast.Subscript(value=value_node,
slice=slice_index_node,
ctx=gast.Load())
self.name_to_var_shape[target_id] = sub_node
has_updated = True
return has_updated
else:
target_id = ast_to_source_code(target_node).strip()
if isinstance(value_node, gast.Name):
if value_node.id in self.name_to_var_shape:
self.name_to_var_shape[target_id] = self.name_to_var_shape[
value_node.id]
return True
if isinstance(value_node, gast.Attribute):
if self.is_var_shape(value_node): # eg: x.shape
self.name_to_var_shape[target_id] = value_node
return True
if isinstance(value_node, gast.Subscript):
if isinstance(value_node.value, gast.Attribute):
if self.is_var_shape(value_node.value): # eg: x.shape[0]
self.name_to_var_shape[target_id] = value_node
return True
return False
def visit_Assign(self, node):
self.src = quoting.unquote(node)
self.mark(node)
self.trivializing = True
self.namer.target = node.targets[0]
if isinstance(node.targets[0], (gast.Subscript, gast.Attribute)):
node.value = self.trivialize(node.value)
node.targets[0] = self.visit(node.targets[0])
elif isinstance(node.targets[0], gast.Tuple):
node.value = self.visit(node.value)
name = self.namer.name(node.targets[0])
target = gast.Name(id=name, ctx=gast.Store(), annotation=None)
for i, elt in enumerate(node.targets[0].elts):
stmt = gast.Assign(targets=[elt],
value=gast.Subscript(
value=gast.Name(id=name,
ctx=gast.Load(),
annotation=None),
slice=gast.Index(value=gast.Num(n=i)),
ctx=gast.Load()))
self.mark(stmt)
self.append(stmt)
node.targets[0] = target
elif not isinstance(node.targets[0], gast.Name):
raise ValueError
node = self.generic_visit(node)
self.namer.target = None
self.trivializing = False
return node
def visit_For(self, node):
target = node.target
if isinstance(target, ast.Tuple) or isinstance(target, ast.List):
renamings = OrderedDict()
self.traverse_tuples(target, (), renamings)
if renamings:
gtarget = self.get_new_id()
node.target = ast.Name(gtarget, node.target.ctx, None)
for rename, state in renamings.items():
nnode = reduce(
lambda x, y: ast.Subscript(
x,
ast.Index(ast.Num(y)),
ast.Load()),
state,
ast.Name(gtarget, ast.Load(), None))
if isinstance(rename, str):
node.body.insert(0,
ast.Assign(
[ast.Name(rename,
ast.Store(),
None)],
nnode)
)
else:
node.body.insert(0, ast.Assign([rename], nnode))
self.generic_visit(node)
return node
def visit_Assign(self, node):
self.generic_visit(node)
# if the rhs is an identifier, we don't need to duplicate it
# otherwise, better duplicate it...
no_tmp = isinstance(node.value, ast.Name)
extra_assign = [] if no_tmp else [node]
for i, t in enumerate(node.targets):
if isinstance(t, ast.Tuple) or isinstance(t, ast.List):
renamings = OrderedDict()
self.traverse_tuples(t, (), renamings)
if renamings:
gtarget = node.value.id if no_tmp else self.get_new_id()
node.targets[i] = ast.Name(gtarget,
node.targets[i].ctx,
None)
for rename, state in renamings.items():
nnode = reduce(
lambda x, y: ast.Subscript(
x,
ast.Index(ast.Num(y)),
ast.Load()),
state,
ast.Name(gtarget, ast.Load(), None))
if isinstance(rename, str):
extra_assign.append(
ast.Assign(
[ast.Name(rename, ast.Store(), None)],
nnode))
else:
extra_assign.append(ast.Assign([rename], nnode))
return extra_assign or node
def _process_variable_assignment(self, source, targets):
if isinstance(source, gast.Call):
func = source.func
if anno.hasanno(func, 'live_val'):
func_obj = anno.getanno(func, 'live_val')
if tf_inspect.isclass(func_obj):
anno.setanno(source, 'is_constructor', True)
anno.setanno(source, 'type', func_obj)
anno.setanno(source, 'type_fqn', anno.getanno(func, 'fqn'))
# TODO(mdan): Raise an error if constructor has side effects.
# We can have a whitelist of no-side-effects constructors.
# We can also step inside the constructor and further analyze.
for t in targets:
if isinstance(t, gast.Tuple):
for i, e in enumerate(t.elts):
self.scope.setval(
anno.getanno(e, anno.Basic.QN),
gast.Subscript(source, gast.Index(i),
ctx=gast.Store()))
elif isinstance(t, (gast.Name, gast.Attribute)):
self.scope.setval(anno.getanno(t, anno.Basic.QN), source)
else:
raise ValueError('Dont know how to handle assignment to %s' %
t)
def visit_Assign(self, node):
self.generic_visit(node)
if isinstance(node.value, gast.Call):
target = node.value.func
if anno.hasanno(target, 'live_val'):
target_obj = anno.getanno(target, 'live_val')
if tf_inspect.isclass(target_obj):
# This is then a constructor.
anno.setanno(node.value, 'type', target_obj)
anno.setanno(node.value, 'type_fqn',
anno.getanno(target, 'fqn'))
# TODO (mdan): Raise an error if constructor has side effects. id:2153 gh:2154
# We can have a whitelist of no-side-effects constructors.
# We can also step inside the constructor and further analyze.
for n in node.targets:
if isinstance(n, gast.Tuple):
for i, e in enumerate(n.elts):
self.scope.setval(
e.id,
gast.Subscript(node.value,
gast.Index(i),
ctx=gast.Store()))
else:
self.scope.setval(n.id, node.value)
return node
def ast(self):
"""AST representation."""
# The caller must adjust the context appropriately.
if self.has_subscript():
return gast.Subscript(value=self.parent.ast(),
slice=gast.Index(self.qn[-1].ast()),
ctx=CallerMustSetThis)
if self.has_attr():
return gast.Attribute(value=self.parent.ast(),
attr=self.qn[-1],
ctx=CallerMustSetThis)
base = self.qn[0]
if isinstance(base, str):
return gast.Name(base,
ctx=CallerMustSetThis,
annotation=None,
type_comment=None)
elif isinstance(base, StringLiteral):
return gast.Constant(base.value, kind=None)
elif isinstance(base, NumberLiteral):
return gast.Constant(base.value, kind=None)
else:
assert False, ('the constructor should prevent types other than '
'str, StringLiteral and NumberLiteral')
def _process_tuple_assignment(self, source, t):
for i, e in enumerate(t.elts):
if isinstance(e, gast.Tuple):
self._process_tuple_assignment(source, e)
else:
self.scope.setval(
anno.getanno(e, anno.Basic.QN),
gast.Subscript(source, gast.Index(i), ctx=gast.Store()))
def _build_var_slice_node(self):
return gast.Subscript(
value=self.iter_node,
slice=gast.Index(value=gast.Name(
id=self.iter_idx_name,
ctx=gast.Load(),
annotation=None,
type_comment=None)),
ctx=gast.Load())
def visit_Name(self, node):
if node.id in self.renamings:
nnode = reduce(
lambda x, y: ast.Subscript(x, ast.Index(ast.Num(y)), ast.Load(
)), self.renamings[node.id],
ast.Name(self.tuple_id, ast.Load(), None))
nnode.ctx = node.ctx
return nnode
return node
def slice2list(self):
# TODO(hamaji): Use 2**63-1 instead.
int_max = 2**31 - 1
if isinstance(self, gast.Slice):
assert self.step is None
def f(x, v):
if x is None:
return Value(np.array([v])).to_tensor(env)
x = eval_ast(x, env)
if x.is_tensor():
return unsqueeze(x.value)
else:
return Value(np.array([x.value])).to_tensor(env)
lower = f(self.lower, 0)
upper = f(self.upper, int_max)
squeeze = [False]
elif isinstance(self, gast.Index):
idx = eval_ast(self.value, env)
if isinstance(idx.value, tuple): # ここにTupleが来うる
# TODO(satos) もっとうまくやったほうがいいかも
vs = [
gast.Index(gast.NameConstant(value=v)) for v in idx.value
]
lower, upper, squeeze = slice2list(gast.ExtSlice(dims=vs))
elif not idx.is_py:
lower = unsqueeze(idx.value)
ot = totensor(1, env)
upper = env.calc(
"Add",
inputs=[idx.to_tensor(env).name, ot.name],
)
upper = unsqueeze(upper)
squeeze = [True]
else:
lower = totensor(np.array([idx.value]), env)
upper_value = idx.value + 1 if idx.value != -1 else int_max
upper = totensor(np.array([upper_value]), env)
squeeze = [True]
elif isinstance(self, gast.ExtSlice):
ds = list(map(slice2list, self.dims))
lower = _concat(
tuple(map(lambda x: castto(x[0], TensorProto.INT64, env), ds)),
0, env)
upper = _concat(
tuple(map(lambda x: castto(x[1], TensorProto.INT64, env), ds)),
0, env)
squeeze = sum(map(lambda x: x[2], ds), [])
else:
raise Exception(self, " is not Python slice")
return lower, upper, squeeze
def visit_Subscript(self, node):
new_slice = self._visit(node.slice)
if isinstance(node.slice, ast.Ellipsis):
new_slice = gast.Index(new_slice)
new_node = gast.Subscript(
self._visit(node.value),
new_slice,
self._visit(node.ctx),
)
gast.copy_location(new_node, node)
new_node.end_lineno = new_node.end_col_offset = None
return new_node
def _build_assign_var_slice_node(self):
var_slice_node = gast.Subscript(
value=self.iter_node,
slice=gast.Index(value=gast.Name(id=self.iter_idx_name,
ctx=gast.Load(),
annotation=None,
type_comment=None)),
ctx=gast.Load(),
)
new_iter_var_name = unique_name.generate(FOR_ITER_VAR_NAME_PREFIX)
target_node, assign_node = create_assign_node(new_iter_var_name,
var_slice_node)
return target_node, assign_node
def apply_to_single_assignments(self, targets, values, apply_fn):
"""Applies a function to each individual assignment.
This function can process a possibly-unpacked (e.g. a, b = c, d) assignment.
It tries to break down the unpacking if possible. In effect, it has the same
effect as passing the assigned values in SSA form to apply_fn.
Examples:
The following will result in apply_fn(a, c), apply_fn(b, d):
a, b = c, d
The following will result in apply_fn(a, c[0]), apply_fn(b, c[1]):
a, b = c
The following will result in apply_fn(a, (b, c)):
a = b, c
It uses the visitor pattern to allow subclasses to process single
assignments individually.
Args:
targets: list, tuple of or individual AST node. Should be used with the
targets field of an ast.Assign node.
values: an AST node.
apply_fn: a function of a single argument, which will be called with the
respective nodes of each single assignment. The signature is
apply_fn(target, value), no return value.
"""
if not isinstance(targets, (list, tuple)):
targets = (targets, )
for target in targets:
if isinstance(target, (gast.Tuple, gast.List)):
for i in range(len(target.elts)):
target_el = target.elts[i]
if isinstance(values, (gast.Tuple, gast.List)):
value_el = values.elts[i]
else:
value_el = gast.Subscript(values,
gast.Index(i),
ctx=gast.Store())
self.apply_to_single_assignments(target_el, value_el,
apply_fn)
else:
# TODO(mdan): Look into allowing to rewrite the AST here.
apply_fn(target, values)
def _build_compare_node(self):
if self.is_for_range_iter():
compare_node = self.iter_args[
0] if self.args_length == 1 else self.iter_args[1]
else:
compare_node = gast.Subscript(
value=gast.Name(
id=self.iter_var_shape_name,
ctx=gast.Load(),
annotation=None,
type_comment=None),
slice=gast.Index(value=gast.Constant(
value=0, kind=None)),
ctx=gast.Load())
return compare_node
def visit_ExtSlice(self, node):
has_ellipsis = any(isinstance(d, ast.Ellipsis) for d in node.dims)
has_slice = any(isinstance(d, ast.Slice) for d in node.dims)
new_dims = self._visit(node.dims)
if has_ellipsis and not has_slice:
new_dims = [
nd.value if isinstance(nd, gast.Index) else nd
for nd in new_dims
]
new_node = gast.Index(gast.Tuple(new_dims, gast.Load()))
else:
new_node = gast.ExtSlice(new_dims)
gast.copy_location(new_node, node)
new_node.end_lineno = new_node.end_col_offset = None
return new_node
def apply_to_single_assignments(targets, values, apply_fn):
"""Applies a function to each individual assignment.
This function can process a possibly-unpacked (e.g. a, b = c, d) assignment.
It tries to break down the unpacking if possible. In effect, it has the same
effect as passing the assigned values in SSA form to apply_fn.
Examples:
The following will result in apply_fn(a, c), apply_fn(b, d):
a, b = c, d
The following will result in apply_fn(a, c[0]), apply_fn(b, c[1]):
a, b = c
The following will result in apply_fn(a, (b, c)):
a = b, c
It uses the visitor pattern to allow subclasses to process single
assignments individually.
Args:
targets: Union[List[ast.AST, ...], Tuple[ast.AST, ...], ast.AST, should be
used with the targets field of an ast.Assign node
values: ast.AST
apply_fn: Callable[[ast.AST, ast.AST], None], called with the respective
nodes of each single assignment
"""
if not isinstance(targets, (list, tuple)):
targets = (targets, )
for target in targets:
if isinstance(target, (gast.Tuple, gast.List)):
for i in range(len(target.elts)):
target_el = target.elts[i]
if isinstance(values, (gast.Tuple, gast.List)):
value_el = values.elts[i]
else:
idx = parsing.parse_expression(str(i))
value_el = gast.Subscript(values,
gast.Index(idx),
ctx=gast.Load())
apply_to_single_assignments(target_el, value_el, apply_fn)
else:
apply_fn(target, values)
def ast(self):
# The caller must adjust the context appropriately.
if self.has_subscript():
return gast.Subscript(self.parent.ast(),
gast.Index(self.qn[-1].ast()), None)
if self.has_attr():
return gast.Attribute(self.parent.ast(), self.qn[-1], None)
base = self.qn[0]
if isinstance(base, str):
return gast.Name(base, None, None)
elif isinstance(base, StringLiteral):
return gast.Str(base.value)
elif isinstance(base, NumberLiteral):
return gast.Num(base.value)
else:
assert False, ('the constructor should prevent types other than '
'str, StringLiteral and NumberLiteral')
def generic_visit(self, node):
if node in self.constant_expressions:
try:
fake_node = ast.Expression(
node.value if isinstance(node, ast.Index) else node)
code = compile(ast.gast_to_ast(fake_node),
'<constant folding>', 'eval')
value = eval(code, self.env)
new_node = to_ast(value)
if(isinstance(node, ast.Index) and
not isinstance(new_node, ast.Index)):
new_node = ast.Index(new_node)
try:
if not ASTMatcher(node).search(new_node):
self.update = True
return new_node
except DamnTooLongPattern as e:
print("W: ", e, " Assume no update happened.")
return Transformation.generic_visit(self, node)
except ConversionError as e:
print('error in constant folding: ', e)
raise
except ToNotEval:
return Transformation.generic_visit(self, node)
except AttributeError as e:
# FIXME union_ function is not handle by constant folding
if "union_" in e.args[0]:
return Transformation.generic_visit(self, node)
elif "pythran" in e.args[0]:
# FIXME: Can be fix giving a Python implementation for
# these functions.
return Transformation.generic_visit(self, node)
raise
except NameError as e:
# FIXME dispatched function are not processed by constant
# folding
if "__dispatch__" in e.args[0]:
return Transformation.generic_visit(self, node)
raise
except Exception as e:
raise PythranSyntaxError(str(e), node)
else:
return Transformation.generic_visit(self, node)
def _process_variable_assignment(self, source, targets):
if isinstance(source, gast.Call):
func = source.func
if anno.hasanno(func, 'live_val'):
func_obj = anno.getanno(func, 'live_val')
if tf_inspect.isclass(func_obj):
# This is then a constructor.
anno.setanno(source, 'type', func_obj)
anno.setanno(source, 'type_fqn', anno.getanno(func, 'fqn'))
# TODO(mdan): Raise an error if constructor has side effects.
# We can have a whitelist of no-side-effects constructors.
# We can also step inside the constructor and further analyze.
for t in targets:
if isinstance(t, gast.Tuple):
for i, e in enumerate(t.elts):
self.scope.setval(e.id,
gast.Subscript(
source, gast.Index(i), ctx=gast.Store()))
else:
self.scope.setval(t.id, source)
def _process_variable_assignment(self, source, targets):
# Special case: constructors.
if isinstance(source, gast.Call):
func = source.func
if anno.hasanno(func, 'live_val'):
func_obj = anno.getanno(func, 'live_val')
if tf_inspect.isclass(func_obj):
anno.setanno(source, 'is_constructor', True)
anno.setanno(source, 'type', func_obj)
anno.setanno(source, 'type_fqn', anno.getanno(func, 'fqn'))
# TODO (mdan): Raise an error if constructor has side effects. id:554
# https://github.com/imdone/tensorflow/issues/555
# We can have a whitelist of no-side-effects constructors.
# We can also step inside the constructor and further analyze.
# Multiple targets mean multiple assignment.
for target in targets:
# Tuple target means unpacking.
if isinstance(target, (gast.Tuple, gast.List)):
for i, target_item in enumerate(target.elts):
# Two cases here:
# 1. Static unpacking, e.g. a, b = c, d
# 2. Dynamic unpacking, e.g. a, b = c
# The former case is optimized away.
if isinstance(source, (gast.Tuple, gast.List)):
source_item = source.elts[i]
else:
source_item = gast.Subscript(source,
gast.Index(i),
ctx=None)
self._process_variable_assignment(source_item,
(target_item, ))
elif isinstance(target, (gast.Name, gast.Attribute)):
target_symbol = anno.getanno(target, anno.Basic.QN)
self.scope.setval(target_symbol, source)
else:
raise ValueError('assignment target has unknown type: %s' %
target)
def visit_Subscript(self, node):
if self.trivializing:
node.value = self.trivialize(node.value)
if isinstance(node.slice, gast.Index):
node.slice.value = self.trivialize(node.slice.value)
elif isinstance(node.slice, gast.Slice):
name = self.namer.name(node.slice)
target = gast.Name(id=name, ctx=gast.Store(), annotation=None)
stmt = gast.Assign(targets=[target], value=None)
self.prepend(stmt)
stmt.value = gast.Call(
func=gast.Name(id='slice', ctx=gast.Load(), annotation=None),
args=[
self.trivialize(arg) if arg else
gast.Name(id='None', ctx=gast.Load(), annotation=None)
for arg in [node.slice.lower, node.slice.upper,
node.slice.step]],
keywords=[])
node.slice = gast.Index(value=gast.Name(id=name, ctx=gast.Load(),
annotation=None))
else:
raise ValueError
return node
MethodIntr(),
"keys":
MethodIntr(),
"pop":
MethodIntr(),
"popitem":
MethodIntr(),
"setdefault":
MethodIntr(lambda self, node: len(node.args) == 3 and self.combine(
node.args[0],
node.args[1],
unary_op=lambda x: cxxtypes.DictType(x, self.result[node.args[2]]),
register=True,
aliasing_type=True),
return_alias=lambda args: {
ast.Subscript(args[0], ast.Index(args[1]), ast.Load())
}.union({args[2]} if len(args) == 3 else set())),
"update":
MethodIntr(update_effects),
"values":
MethodIntr(),
"viewitems":
MethodIntr(),
"viewkeys":
MethodIntr(),
"viewvalues":
MethodIntr(),
},
"file": {
# Member variables
"closed": AttributeIntr(return_type=NamedType("bool")),
def visit_Subscript(self, node):
if self.trivializing:
node.value = self.trivialize(node.value)
node.slice = gast.Index(value=self.trivialize_slice(node.slice))
return node
def visit_Index(self, node):
value = self.visit(node.value)
if value is not node.value:
return ast.Index(value)
else:
return node
