Skip to content

Commit d131540

Browse files
committed
Fix segfault bug in Alloc C implementation
1 parent 81e7dff commit d131540

File tree

2 files changed

+17
-3
lines changed

2 files changed

+17
-3
lines changed

pytensor/tensor/basic.py

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -1506,7 +1506,7 @@ def c_code(self, node, name, inp, out, sub):
15061506
if (need_new_out)
15071507
{{
15081508
Py_XDECREF({zz});
1509-
{zz} = (PyArrayObject*) PyArray_SimpleNew({ndim}, shape, PyArray_TYPE((PyArrayObject*) py_{vv}));
1509+
{zz} = (PyArrayObject*) PyArray_SimpleNew({ndim}, shape, PyArray_TYPE({vv}));
15101510
if (!{zz})
15111511
{{
15121512
PyErr_SetString(PyExc_MemoryError, "alloc failed");
@@ -1522,7 +1522,7 @@ def c_code(self, node, name, inp, out, sub):
15221522
return code
15231523

15241524
def c_code_cache_version(self):
1525-
return (2,)
1525+
return (3,)
15261526

15271527
def infer_shape(self, fgraph, node, input_shapes):
15281528
return [node.inputs[1:]]

tests/tensor/test_basic.py

Lines changed: 15 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -14,7 +14,7 @@
1414
from pytensor.compile.mode import Mode, get_default_mode
1515
from pytensor.compile.ops import DeepCopyOp
1616
from pytensor.gradient import grad, hessian
17-
from pytensor.graph.basic import Apply
17+
from pytensor.graph.basic import Apply, equal_computations
1818
from pytensor.graph.op import Op
1919
from pytensor.graph.replace import clone_replace
2020
from pytensor.misc.safe_asarray import _asarray
@@ -719,6 +719,7 @@ class TestAlloc:
719719
shared = staticmethod(pytensor.shared)
720720
allocs = [Alloc()] * 3
721721

722+
722723
def setup_method(self):
723724
self.rng = np.random.default_rng(seed=utt.fetch_seed())
724725

@@ -850,6 +851,19 @@ def test_static_shape(self):
850851
with pytest.raises(ValueError, match=msg):
851852
at.alloc(x, 3, 1, 6)
852853

854+
def test_alloc_of_view_linker(self):
855+
"""Check we can allocate a new array properly in the C linker when input is a view."""
856+
x_v = vector("x", shape=(None,))
857+
dim_len = scalar("dim_len", dtype=int)
858+
out = alloc(specify_shape(x_v, (1,)), 5, dim_len)
859+
860+
f = pytensor.function([x_v, dim_len], out, mode=Mode("c"))
861+
assert equal_computations(
862+
f.maker.fgraph.outputs, [alloc(specify_shape(x_v, (1,)), 5, dim_len)]
863+
)
864+
865+
np.testing.assert_array_equal(f(x=np.zeros((1,)), dim_len=3), np.zeros((5, 3)))
866+
853867

854868
def test_infer_shape():
855869
with pytest.raises(TypeError, match="^Shapes must be scalar integers.*"):

0 commit comments

Comments
 (0)