|
14 | 14 | from pytensor.compile.mode import Mode, get_default_mode
|
15 | 15 | from pytensor.compile.ops import DeepCopyOp
|
16 | 16 | from pytensor.gradient import grad, hessian
|
17 |
| -from pytensor.graph.basic import Apply |
| 17 | +from pytensor.graph.basic import Apply, equal_computations |
18 | 18 | from pytensor.graph.op import Op
|
19 | 19 | from pytensor.graph.replace import clone_replace
|
20 | 20 | from pytensor.misc.safe_asarray import _asarray
|
@@ -719,6 +719,7 @@ class TestAlloc:
|
719 | 719 | shared = staticmethod(pytensor.shared)
|
720 | 720 | allocs = [Alloc()] * 3
|
721 | 721 |
|
| 722 | + |
722 | 723 | def setup_method(self):
|
723 | 724 | self.rng = np.random.default_rng(seed=utt.fetch_seed())
|
724 | 725 |
|
@@ -850,6 +851,19 @@ def test_static_shape(self):
|
850 | 851 | with pytest.raises(ValueError, match=msg):
|
851 | 852 | at.alloc(x, 3, 1, 6)
|
852 | 853 |
|
| 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 | + |
853 | 867 |
|
854 | 868 | def test_infer_shape():
|
855 | 869 | with pytest.raises(TypeError, match="^Shapes must be scalar integers.*"):
|
|
0 commit comments