提交 246714b6 authored 作者: ricardoV94's avatar ricardoV94 提交者: Ricardo Vieira

Fix bug in Numba inplace vectorize code with multiple outputs

上级 477fbafb
......@@ -265,7 +265,7 @@ def _vectorized(
ctx.nrt.incref(
builder,
sig.return_type.types[inplace_idx],
outputs[inplace_idx]._get_value(),
outputs[inplace_idx]._getvalue(),
)
return ctx.make_tuple(
builder, sig.return_type, [out._getvalue() for out in outputs]
......
......@@ -12,7 +12,7 @@ from pytensor import config, function
from pytensor.compile import get_mode
from pytensor.compile.ops import deep_copy_op
from pytensor.gradient import grad
from pytensor.scalar import float64
from pytensor.scalar import Composite, float64
from pytensor.tensor.elemwise import CAReduce, DimShuffle, Elemwise
from pytensor.tensor.math import All, Any, Max, Min, Prod, ProdWithoutZeros, Sum
from pytensor.tensor.special import LogSoftmax, Softmax, SoftmaxGrad
......@@ -548,7 +548,7 @@ def test_Argmax(x, axes, exc):
)
def test_elemwise_out_type():
def test_elemwise_inplace_out_type():
# Create a graph with an elemwise
# Ravel failes if the elemwise output type is reported incorrectly
x = pt.matrix()
......@@ -563,6 +563,28 @@ def test_elemwise_out_type():
assert func(x_val).shape == (18,)
def test_elemwise_multiple_inplace_outs():
x = pt.vector()
y = pt.vector()
x_ = pt.scalar_from_tensor(x[0])
y_ = pt.scalar_from_tensor(y[0])
out_ = x_ + 1, y_ + 1
composite_op = Composite([x_, y_], out_)
elemwise_op = Elemwise(composite_op, inplace_pattern={0: 0, 1: 1})
out = elemwise_op(x, y)
fn = function([x, y], out, mode="NUMBA", accept_inplace=True)
x_test = np.array([1, 2, 3], dtype=config.floatX)
y_test = np.array([4, 5, 6], dtype=config.floatX)
out1, out2 = fn(x_test, y_test)
assert out1 is x_test
assert out2 is y_test
np.testing.assert_allclose(out1, [2, 3, 4])
np.testing.assert_allclose(out2, [5, 6, 7])
def test_scalar_loop():
a = float64("a")
scalar_loop = pytensor.scalar.ScalarLoop([a], [a + a])
......
Markdown 格式
0%
您添加了 0 到此讨论。请谨慎行事。
请先完成此评论的编辑!
注册 或者 后发表评论