def make_py_thunk(
self,
node: Apply,
storage_map: StorageMapType,
compute_map: ComputeMapType,
no_recycling: List[Variable],
debug: bool = False,
) -> ThunkType:
"""Make a Python thunk.
Like :meth:`Op.make_thunk` but only makes Python thunks.
"""
node_input_storage = [storage_map[r] for r in node.inputs]
node_output_storage = [storage_map[r] for r in node.outputs]
if debug and hasattr(self, "debug_perform"):
p = node.op.debug_perform # type: ignore
else:
p = node.op.perform
params = node.run_params()
if params is NoParams:
# default arguments are stored in the closure of `rval`
@is_thunk_type
def rval(
p=p, i=node_input_storage, o=node_output_storage, n=node, params=None
):
r = p(n, [x[0] for x in i], o)
for o in node.outputs:
compute_map[o][0] = True
return r
else:
params_val = node.params_type.filter(params)
@is_thunk_type
def rval(
p=p,
i=node_input_storage,
o=node_output_storage,
n=node,
params=params_val,
):
r = p(n, [x[0] for x in i], o, params)
for o in node.outputs:
compute_map[o][0] = True
return r
rval.inputs = node_input_storage
rval.outputs = node_output_storage
setattr(rval, "perform", p)
rval.lazy = False
return rval
def make_py_thunk(
self,
node: Apply,
storage_map: StorageMapType,
compute_map: ComputeMapType,
no_recycling: bool,
debug: bool = False,
) -> ThunkType:
"""Make a Python thunk.
Like `Op.make_thunk` but only makes python thunks.
"""
node_input_storage = [storage_map[r] for r in node.inputs]
node_output_storage = [storage_map[r] for r in node.outputs]
if debug:
p = node.op.debug_perform
else:
p = node.op.perform
params = node.run_params()
if params is NoParams:
# default arguments are stored in the closure of `rval`
@rval_decorator
def rval(p=p, i=node_input_storage, o=node_output_storage, n=node):
r = p(n, [x[0] for x in i], o)
for o in node.outputs:
compute_map[o][0] = True
return r
else:
params_val = node.params_type.filter(params)
@rval_decorator
def rval(
p=p,
i=node_input_storage,
o=node_output_storage,
n=node,
params=params_val,
):
r = p(n, [x[0] for x in i], o, params)
for o in node.outputs:
compute_map[o][0] = True
return r
rval.inputs = node_input_storage
rval.outputs = node_output_storage
rval.perform = p
rval.lazy = False
return rval