def __repr__(self):
inputs = self.get_sanitized_inputs()
if self.at_idx:
return self.fmt_with_idx.format(name=self.name,
in1=self.input1,
idx=', '.join(
[str(i) for i in self.at_idx]),
in2=self.input2)
elif self.slices:
lhs_eval = get_unique_name(inputs[0] + '_eval')
rhs_eval = get_unique_name(inputs[1] + '_eval')
return self.slice_fmt.format(
name=self.name,
lhs_eval=lhs_eval,
rhs_eval=rhs_eval,
lhs=inputs[0],
rhs=inputs[1],
start='{}, {}'.format(self.slice_tuples[0][0],
self.slice_tuples[1][0]),
end='{}, {}'.format(self.slice_tuples[0][1],
self.slice_tuples[1][1]))
return self.fmt.format(in1=self.input1,
in2=self.input2,
tranpose=self.transpose)
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,
value=None,
name='array',
shape=(),
dtype=config.DTYPE,
producer=None,
**kwargs):
if value is None:
self.value = numpy.array(None)
self.shape = shape
else:
self.value = numpy.array(value, dtype=object)
self.shape = self.value.shape
if len(shape) > 2:
raise NotImplementedError('Only up to 2 dimensions supported.')
value = expander.check_type(*self.value.flatten().tolist())
self.value = numpy.array(value, dtype=object).reshape(self.shape)
op = expander.CreateArrayOp(value, self, self.shape)
producer = op
self.ndim = len(shape)
self.dtype = numpy.dtype(dtype)
self.name = utils.get_unique_name(name)
self.producer = producer
self.assignment = []
self.kwargs = kwargs
def __init__(self, out, args=None, fun_name='func', file_name=None):
"""
Construct a class builder
Args:
graph (nx.DiGraph): The jet graph containing all Ops
out (list): List of nodes which form the output of the compiled
function
fun_name (str): Name suffix of the compiled class.
If fun_name = FunName, the compiled class
would be called JetFunName.
"""
if file_name is None:
file_name = get_unique_name('jet_autogenerated')
self.graph = expander.graph
self.args = [] if args is None else [arg.producer for arg in args]
self.ops = []
self.variables = []
self.constants = []
self.Op = OpCollector()
self.file_name = file_name
self.fun_name = fun_name + 'Func'
self.class_name = fun_name.title() + 'Class'
self.extract_return(out)
# extract sequential list, starting at output
current_nodes = self._extract_nodes_for_return(out)
self.subgraph = self.graph.subgraph(current_nodes)
topo_sorted = nx.topological_sort(self.subgraph)
if not config.debug and config.merge:
self._merge_ops(self.subgraph, topo_sorted)
for el in topo_sorted:
if el not in current_nodes:
continue
if el.op == 'Placeholder':
if args is None:
self.args.append(el)
elif el.op == 'Variable':
self.variables.append(el)
elif el.op == 'Const':
self.constants.append(el)
else:
self.ops.append(el)
self.args = sorted(self.args, key=lambda x: x.placeholder_count)
if config.draw_graph:
import burn
burn.draw(self.subgraph, outputs=out, name=fun_name)
if config.draw_graph_raw:
import burn
burn.draw(self.graph, outputs=out, name=fun_name + '_raw')
def add_to_graph(self):
if hasattr(self, 'name'): # TODO dominique: is this check here? the class has a property 'name'
graph.add_node(self, name=self.name)
else:
name=get_unique_name('uniquename')
graph.add_node(self, name=name)
if self.inputs is not None:
for arr in self.inputs:
if hasattr(arr, 'assignment') and arr.assignment:
graph.add_edge(arr.assignment[-1], self, array=arr.name)
elif hasattr(arr, 'producer'):
graph.add_edge(arr.producer, self, array=arr.name)