def __init__(self, inputs, *args, **kwargs):
self.inputs_raw = inputs
self.inputs = check_type(*inputs)
self.output = None
self.init_op(self.inputs, *args, **kwargs)
if config.debug or config.group_class or config.group_func:
self.caller_info = get_caller_info('expander.py', 'intake.py')
else:
self.caller_info = None
self.add_to_graph()
def __init__(self, inputs, producer, shape):
self.op = 'CreateArray'
self.inputs = inputs
self.nested_input = np.array(inputs, dtype=object).reshape(shape)
self.output = producer # TODO dominique: this is some weird naming convention: a producer is an output? isn't a producer normally responsible for the input?
self.shape = shape
if config.debug or config.group_class or config.group_func:
self.caller_info = get_caller_info('expander.py', 'intake.py')
else:
self.caller_info = None
self.add_to_graph()
def wrapper(*args):
if not jet.jet_mode:
return func(*args)
func_id = id(func)
func_cached = _func_cached_dict[func_id]['func']
if func_cached is not None:
return func_cached(*args)
shapes = _func_cached_dict[func_id]['shapes']
if inspect.ismethod(func):
arg_names = func.__code__.co_varnames[1:func.__code__.
co_argcount]
else:
arg_names = func.__code__.co_varnames[:func.__code__.
co_argcount]
if len(arg_names) != len(args):
assert (len(arg_names) == 0)
arg_names = [get_unique_name('ph') for each in args]
if len(shapes) != len(arg_names) and shapes:
raise ValueError(
'Shapes length does not match the arguments length.')
if not shapes:
shapes = [
arg.shape if hasattr(arg, 'shape') else () for arg in args
]
_func_cached_dict[func_id]['shapes'] = shapes
ph = [
placeholder(name=arg[1], shape=shapes[arg[0]])
for arg in enumerate(arg_names)
]
fun_name = func.__code__.co_name
if fun_name == '<lambda>':
fun_name = get_unique_name('lambda')
jb = JetBuilder(args=ph,
out=func(*ph),
file_name=get_unique_name(
sanitize_name('{}_{}_{func_name}'.format(
*get_caller_info('jit.py')[1:-1],
func_name=fun_name))),
fun_name=get_unique_name(fun_name))
jet_class = getattr(jb.build(), jb.class_name)
jet_func = getattr(jet_class(), jb.fun_name)
_func_cached_dict[func_id]['func'] = jet_func
return jet_func(*args)
def __init__(self, inputs, at_idx=None, slices=None):
self.op = 'Assign'
self.inputs = check_type(*inputs)
self.slices = slices
if slices is not None:
if len(slices) == 1:
slice_shape = np.zeros(inputs[0].shape)[slices[0]].shape
else:
slice_shape = np.zeros(inputs[0].shape)[slices[0],
slices[1]].shape
self.output = intake.array(name=self.op,
dtype=inputs[0].dtype,
shape=inputs[0].shape,
producer=self,
slice_shape=slice_shape)
self.__repr__ = lambda: '{}[{}, {}]'.format(
self.inputs[0].name, slice_to_str(slices[0]),
slice_to_str(slices[1])) + '\\n' + str(self.output.shape)
else:
self.output = intake.array(name=self.op,
dtype=inputs[0].dtype,
shape=inputs[0].shape,
producer=self)
self.at_idx = at_idx
if at_idx:
if len(at_idx) == 1:
self.__repr__ = lambda: '{}[{}]'.format(
self.inputs[0].name, at_idx[0]) + '\\n' + str(self.output.
shape)
else:
self.__repr__ = lambda: '{}[{}, {}]'.format(
self.inputs[0].name, at_idx[0], at_idx[1]) + '\\n' + str(
self.output.shape)
# special thing for assign operator... have to add the others
# as dependency
successors = graph.successors(inputs[0].last_producer)
for s in successors:
if s != self.inputs[0].last_producer:
graph.add_edge(s, self, edge_type='helper')
if config.debug or config.group_class or config.group_func:
self.caller_info = get_caller_info('expander.py', 'intake.py')
else:
self.caller_info = None
self.add_to_graph()
self.inputs[0].assignment.append(self)