def visit_IfExp(self, expr):
cond = self.visit(expr.test)
true = self.visit(expr.body)
false = self.visit(expr.orelse)
if is_iterable(true) or is_iterable(false):
make = np.vectorize(lambda c, t, f: sy2.IfExp(c, t, f))
return make(cond, true, false)
else:
return sy2.IfExp(cond, true, false)
def visit_ListComp(self, expr):
assert len(expr.generators) == 1, "Multiple generators not implemented"
gen = expr.generators[0]
iterator = self.visit(gen.iter)
# if is_symbolic(iterator):
# raise NotImplementedError("Symbolic generators not implemented")
if not is_iterable(iterator):
raise ValueError("%r is not iterable" % iterator)
target = gen.target
if isinstance(target, ast.Name):
temps = [target.id]
elif (isinstance(target, ast.Tuple) and
all(isinstance(e, ast.Name) for e in target.elts)):
temps = [e.id for e in target.elts]
else:
raise ValueError("Unrecognized target for list comprehension")
for temp in temps:
assert temp not in self.temps
result = []
for i, x in enumerate(iterator):
if len(temps) == 1:
self.temps[temps[0]] = x
else:
for temp, xj in zip(temps, x):
self.temps[temp] = xj
result.append(self.visit(expr.elt))
self._check_vector_length(i)
for temp in temps:
del self.temps[temp]
return result
def __init__(self, source, globals_dict, closure_dict,
in_dims=None, out_dim=None):
self.source = source
self.globals = globals_dict
self.closures = closure_dict
# self.in_dims = in_dims
self.out_dim = out_dim
self.locals = {}
self.temps = {} # for comprehensions
### parse and make code
a = ast.parse(source)
ff = Function_Finder()
ff.visit(a)
function_def = ff.fn_node
self.arg_names = [arg.id for arg in function_def.args.args]
if in_dims is None:
in_dims = [1] * len(self.arg_names) # assume all scalars
elif not is_iterable(in_dims):
in_dims = (in_dims,)
self.arg_dims = dict(zip(self.arg_names, in_dims))
for v in self.arg_dims.values():
self._check_vector_length(v)
self.args = dict(zip(
self.arg_names,
[self._new_symbolic_vector(name, dims)
for name, dims in zip(self.arg_names, in_dims)]
))
if isinstance(function_def, ast.FunctionDef):
self.function_name = function_def.name
self.body = self.visit_block(function_def.body)
elif isinstance(function_def, ast.Lambda):
if hasattr(function_def, 'targets'):
self.function_name = function_def.targets[0].id
else:
self.function_name = "<lambda>"
r = ast.Return() #wrap lambda expression to look like a one-line function
r.value = function_def.body
r.lineno = 1
r.col_offset = 4
self.body = self.visit_block([r])
else:
raise RuntimeError(
"Expected function definition or lambda function assignment, "
"got " + str(type(function_def)))
def _print_Return(self, expr, indent=0):
print_return = lambda i, v: self._indent(
"%s[%d] = %s;" % (self.OUTPUT, i, self._print(v)), indent)
if is_iterable(expr.value):
