def test_codegen_with_DictOfNamedArrays(ctx_factory): # noqa
ctx = ctx_factory()
queue = cl.CommandQueue(ctx)
namespace = pt.Namespace()
x = Placeholder(namespace, "x", (5, ), np.int)
y = Placeholder(namespace, "y", (5, ), np.int)
x_in = np.array([1, 2, 3, 4, 5])
y_in = np.array([6, 7, 8, 9, 10])
result = pt.DictOfNamedArrays(dict(x_out=x, y_out=y))
# Without return_dict.
prog = pt.generate_loopy(result, target=pt.PyOpenCLTarget(queue))
_, (x_out, y_out) = prog(x=x_in, y=y_in)
assert (x_out == x_in).all()
assert (y_out == y_in).all()
# With return_dict.
prog = pt.generate_loopy(result,
target=pt.PyOpenCLTarget(queue),
options=lp.Options(return_dict=True))
_, outputs = prog(x=x_in, y=y_in)
assert (outputs["x_out"] == x_in).all()
assert (outputs["y_out"] == y_in).all()
def map_data_wrapper(self, expr: DataWrapper) -> Array:
self.bound_arguments[expr.name] = expr.data
return Placeholder(namespace=self.namespace,
name=expr.name,
shape=expr.shape,
dtype=expr.dtype,
tags=expr.tags)
def map_placeholder(self, expr: Placeholder) -> Array:
name = expr.name
if name is None:
name = self.var_name_gen("_pt_in")
return Placeholder(name=name,
shape=tuple(self.rec(s) if isinstance(s, Array) else s
for s in expr.shape),
dtype=expr.dtype,
axes=expr.axes,
tags=expr.tags)
def test_scalar_placeholder(ctx_factory):
ctx = ctx_factory()
queue = cl.CommandQueue(ctx)
namespace = pt.Namespace()
x = Placeholder(namespace, "x", (), np.int)
prog = pt.generate_loopy(x, target=pt.PyOpenCLTarget(queue))
x_in = np.array(1)
_, (x_out, ) = prog(x=x_in)
assert np.array_equal(x_out, x_in)
def test_basic_codegen(ctx_factory):
ctx = ctx_factory()
queue = cl.CommandQueue(ctx)
namespace = pt.Namespace()
x = Placeholder(namespace, "x", (5, ), np.int)
prog = pt.generate_loopy(x * x, target=pt.PyOpenCLTarget(queue))
x_in = np.array([1, 2, 3, 4, 5])
_, (out, ) = prog(x=x_in)
assert (out == x_in * x_in).all()
def map_data_wrapper(self, expr: DataWrapper) -> Array:
name = expr.name
if name is None:
name = self.var_name_gen("_pt_data")
self.bound_arguments[name] = expr.data
return Placeholder(name=name,
shape=tuple(self.rec(s) if isinstance(s, Array) else s
for s in expr.shape),
dtype=expr.dtype,
axes=expr.axes,
tags=expr.tags)
def map_placeholder(self, expr: Placeholder) -> Array:
return Placeholder(namespace=self.namespace,
name=expr.name,
shape=expr.shape,
dtype=expr.dtype,
tags=expr.tags)