提交 d1315406 authored 作者: Ricardo Vieira's avatar Ricardo Vieira 提交者: Ricardo Vieira

Fix segfault bug in Alloc C implementation

上级 81e7dff6
...@@ -1506,7 +1506,7 @@ class Alloc(COp): ...@@ -1506,7 +1506,7 @@ class Alloc(COp):
if (need_new_out) if (need_new_out)
{{ {{
Py_XDECREF({zz}); Py_XDECREF({zz});
{zz} = (PyArrayObject*) PyArray_SimpleNew({ndim}, shape, PyArray_TYPE((PyArrayObject*) py_{vv})); {zz} = (PyArrayObject*) PyArray_SimpleNew({ndim}, shape, PyArray_TYPE({vv}));
if (!{zz}) if (!{zz})
{{ {{
PyErr_SetString(PyExc_MemoryError, "alloc failed"); PyErr_SetString(PyExc_MemoryError, "alloc failed");
...@@ -1522,7 +1522,7 @@ class Alloc(COp): ...@@ -1522,7 +1522,7 @@ class Alloc(COp):
return code return code
def c_code_cache_version(self): def c_code_cache_version(self):
return (2,) return (3,)
def infer_shape(self, fgraph, node, input_shapes): def infer_shape(self, fgraph, node, input_shapes):
return [node.inputs[1:]] return [node.inputs[1:]]
......
...@@ -14,7 +14,7 @@ from pytensor.compile.io import In, Out ...@@ -14,7 +14,7 @@ from pytensor.compile.io import In, Out
from pytensor.compile.mode import Mode, get_default_mode from pytensor.compile.mode import Mode, get_default_mode
from pytensor.compile.ops import DeepCopyOp from pytensor.compile.ops import DeepCopyOp
from pytensor.gradient import grad, hessian from pytensor.gradient import grad, hessian
from pytensor.graph.basic import Apply from pytensor.graph.basic import Apply, equal_computations
from pytensor.graph.op import Op from pytensor.graph.op import Op
from pytensor.graph.replace import clone_replace from pytensor.graph.replace import clone_replace
from pytensor.misc.safe_asarray import _asarray from pytensor.misc.safe_asarray import _asarray
...@@ -719,6 +719,7 @@ class TestAlloc: ...@@ -719,6 +719,7 @@ class TestAlloc:
shared = staticmethod(pytensor.shared) shared = staticmethod(pytensor.shared)
allocs = [Alloc()] * 3 allocs = [Alloc()] * 3
def setup_method(self): def setup_method(self):
self.rng = np.random.default_rng(seed=utt.fetch_seed()) self.rng = np.random.default_rng(seed=utt.fetch_seed())
...@@ -850,6 +851,19 @@ class TestAlloc: ...@@ -850,6 +851,19 @@ class TestAlloc:
with pytest.raises(ValueError, match=msg): with pytest.raises(ValueError, match=msg):
at.alloc(x, 3, 1, 6) at.alloc(x, 3, 1, 6)
def test_alloc_of_view_linker(self):
"""Check we can allocate a new array properly in the C linker when input is a view."""
x_v = vector("x", shape=(None,))
dim_len = scalar("dim_len", dtype=int)
out = alloc(specify_shape(x_v, (1,)), 5, dim_len)
f = pytensor.function([x_v, dim_len], out, mode=Mode("c"))
assert equal_computations(
f.maker.fgraph.outputs, [alloc(specify_shape(x_v, (1,)), 5, dim_len)]
)
np.testing.assert_array_equal(f(x=np.zeros((1,)), dim_len=3), np.zeros((5, 3)))
def test_infer_shape(): def test_infer_shape():
with pytest.raises(TypeError, match="^Shapes must be scalar integers.*"): with pytest.raises(TypeError, match="^Shapes must be scalar integers.*"):
......
Markdown 格式
0%
您添加了 0 到此讨论。请谨慎行事。
请先完成此评论的编辑!
注册 或者 后发表评论