Repository: pymc-devs/pymc

Issue:
Missing static information for jax backend for models with missing values
### Description

The following fails in the jax backend when it tries to allocate a new array like an existing array with unknown shape:

```python
import pymc as pm
import numpy as np
import pytensor
import pytensor.tensor as pt

with pm.Model(coords={"a": [0, 1, 2]}) as model:
#with pm.Model() as model:
with pm.Model(coords={"a": [0, 1, 2]}) as model:
    y = pm.Normal("y", 0, observed=[0, 0, np.nan], dims="a")


from pymc.model.transform.optimization import freeze_dims_and_data

model = freeze_dims_and_data(model, dims=model.coords.keys())
outputs = pytensor.clone_replace(model.unobserved_value_vars, model.rvs_to_values)
func = pytensor.function(model.value_vars, outputs, mode="JAX")
func(np.ones(1))
```

```
File [/tmp/tmp5pyizklf:5](http://localhost:7890/tmp/tmp5pyizklf#line=4), in jax_funcified_fgraph(y_unobserved, a)
      3 tensor_variable = deepcopyop(y_unobserved)
      4 # AllocEmpty{dtype='float64'}(a)
----> 5 tensor_variable_1 = allocempty(a)
      6 # AdvancedSetSubtensor(AllocEmpty{dtype='float64'}.0, y_unobserved, [False False  True])
      7 tensor_variable_2 = advancedincsubtensor(tensor_variable_1, y_unobserved, tensor_constant)

File [~/git/pytensor/pytensor/link/jax/dispatch/tensor_basic.py:39](http://localhost:7890/lab/tree/pymc-labs/pytensor/pytensor/link/jax/dispatch/tensor_basic.py#line=38), in jax_funcify_AllocEmpty.<locals>.allocempty(*shape)
     38 def allocempty(*shape):
---> 39     return jnp.empty(shape, dtype=op.dtype)

File [~/micromamba/envs/dev-cuda/lib/python3.11/site-packages/jax/_src/numpy/lax_numpy.py:2785](http://localhost:7890/home/adr/micromamba/envs/dev-cuda/lib/python3.11/site-packages/jax/_src/numpy/lax_numpy.py#line=2784), in empty(shape, dtype, device)
   2783 if (m := _check_forgot_shape_tuple("empty", shape, dtype)): raise TypeError(m)
   2784 dtypes.check_user_dtype_supported(dtype, "empty")
-> 2785 return zeros(shape, dtype, device=device)

File [~/micromamba/envs/dev-cuda/lib/python3.11/site-packages/jax/_src/numpy/lax_numpy.py:2765](http://localhost:7890/home/adr/micromamba/envs/dev-cuda/lib/python3.11/site-packages/jax/_src/numpy/lax_numpy.py#line=2764), in zeros(shape, dtype, device)
   2763 if (m := _check_forgot_shape_tuple("zeros", shape, dtype)): raise TypeError(m)
   2764 dtypes.check_user_dtype_supported(dtype, "zeros")
-> 2765 shape = canonicalize_shape(shape)
   2766 return lax.full(shape, 0, _jnp_dtype(dtype), sharding=_normalize_to_sharding(device))

File [~/micromamba/envs/dev-cuda/lib/python3.11/site-packages/jax/_src/numpy/lax_numpy.py:100](http://localhost:7890/home/adr/micromamba/envs/dev-cuda/lib/python3.11/site-packages/jax/_src/numpy/lax_numpy.py#line=99), in canonicalize_shape(shape, context)
     98   return core.canonicalize_shape((shape,), context)  # type: ignore
     99 else:
--> 100   return core.canonicalize_shape(shape, context)

File [~/micromamba/envs/dev-cuda/lib/python3.11/site-packages/jax/_src/core.py:1647](http://localhost:7890/home/adr/micromamba/envs/dev-cuda/lib/python3.11/site-packages/jax/_src/core.py#line=1646), in canonicalize_shape(shape, context)
   1645 except TypeError:
   1646   pass
-> 1647 raise _invalid_shape_error(shape, context)

TypeError: Shapes must be 1D sequences of concrete values of integer type, got (Traced<ShapedArray(int64[])>with<DynamicJaxprTrace(level=1[/0](http://localhost:7890/0))>,).
If using `jit`, try using `static_argnums` or applying `jit` to smaller subfunctions.
The error occurred while tracing the function jax_funcified_fgraph at [/tmp/tmp5pyizklf:1](http://localhost:7890/tmp/tmp5pyizklf#line=0) for jit. This concrete value was not available in Python because it depends on the value of the argument a.
Apply node that caused the error: DeepCopyOp(y_unobserved)
Toposort index: 0
Inputs types: [TensorType(float64, shape=(1,))]
Inputs shapes: [(1,), ()]
Inputs strides: [(8,), ()]
Inputs values: [array([1.]), array(3)]
Outputs clients: [['output']]

HINT: Re-running with most PyTensor optimizations disabled could provide a back-trace showing when this node was created. This can be done by setting the PyTensor flag 'optimizer=fast_compile'. If that does not work, PyTensor optimizations can be disabled with 'optimizer=None'.
HINT: Use the PyTensor flag `exception_verbosity=high` for a debug print-out and storage map footprint of this Apply node.
```

The dprint of the function:
```
DeepCopyOp [id A] <Vector(float64, shape=(1,))> 0
 └─ y_unobserved [id B] <Vector(float64, shape=(1,))>
AdvancedSetSubtensor [id C] <Vector(float64, shape=(?,))> 'y' 3
 ├─ AdvancedSetSubtensor [id D] <Vector(float64, shape=(?,))> 2
 │  ├─ AllocEmpty{dtype='float64'} [id E] <Vector(float64, shape=(?,))> 1
 │  │  └─ a [id F] <Scalar(int64, shape=())>
 │  ├─ y_unobserved [id B] <Vector(float64, shape=(1,))>
 │  └─ [False False  True] [id G] <Vector(bool, shape=(3,))>
 ├─ [0. 0.] [id H] <Vector(float64, shape=(2,))>
 └─ [ True  True False] [id I] <Vector(bool, shape=(3,))>
```

I think the problem is that `AllocEmpty` doesn't have a known shape.

Everything works fine if no dimensions are specified:

```python
with pm.Model() as model:
    y = pm.Normal("y", 0, observed=[0, 0, np.nan])
```

In this case the shape is known:
```
DeepCopyOp [id A] <Vector(float64, shape=(1,))> 0
 └─ y_unobserved [id B] <Vector(float64, shape=(1,))>
AdvancedSetSubtensor [id C] <Vector(float64, shape=(?,))> 'y' 3
 ├─ AdvancedSetSubtensor [id D] <Vector(float64, shape=(?,))> 2
 │  ├─ AllocEmpty{dtype='float64'} [id E] <Vector(float64, shape=(3,))> 1
 │  │  └─ 3 [id F] <Scalar(int64, shape=())>
 │  ├─ y_unobserved [id B] <Vector(float64, shape=(1,))>
 │  └─ [False False  True] [id G] <Vector(bool, shape=(3,))>
 ├─ [0. 0.] [id H] <Vector(float64, shape=(2,))>
 └─ [ True  True False] [id I] <Vector(bool, shape=(3,))>
```



Discussion:
codecov[bot]: ## [Codecov](https://app.codecov.io/gh/pymc-devs/pymc/pull/7388?dropdown=coverage&src=pr&el=h1&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=pymc-devs) Report
All modified and coverable lines are covered by tests :white_check_mark:
> Project coverage is 92.19%. Comparing base [(`7af0a87`)](https://app.codecov.io/gh/pymc-devs/pymc/commit/7af0a873550d1c17253acaa58fca54d54993e1c1?dropdown=coverage&el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=pymc-devs) to head [(`e2f86cb`)](https://app.codecov.io/gh/pymc-devs/pymc/commit/e2f86cb3284ca12d7023380f62c66b628ba6ca71?dropdown=coverage&el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=pymc-devs).

<details><summary>Additional details and impacted files</summary>


[![Impacted file tree graph](https://app.codecov.io/gh/pymc-devs/pymc/pull/7388/graphs/tree.svg?width=650&height=150&src=pr&token=JFuXtOJ4Cb&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=pymc-devs)](https://app.codecov.io/gh/pymc-devs/pymc/pull/7388?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=pymc-devs)

```diff
@@           Coverage Diff           @@
##             main    #7388   +/-   ##
=======================================
  Coverage   92.19%   92.19%           
=======================================
  Files         103      103           
  Lines       17214    17214           
=======================================
  Hits        15870    15870           
  Misses       1344     1344           
```

| [Files](https://app.codecov.io/gh/pymc-devs/pymc/pull/7388?dropdown=coverage&src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=pymc-devs) | Coverage Δ | |
|---|---|---|
| [pymc/pytensorf.py](https://app.codecov.io/gh/pymc-devs/pymc/pull/7388?src=pr&el=tree&filepath=pymc%2Fpytensorf.py&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=pymc-devs#diff-cHltYy9weXRlbnNvcmYucHk=) | `90.44% <100.00%> (ø)` | |

</details>

Patch:
diff --git a/pymc/pytensorf.py b/pymc/pytensorf.py
index 2d6910fc4..1ffbb3cd1 100644
--- a/pymc/pytensorf.py
+++ b/pymc/pytensorf.py
@@ -1045,7 +1045,7 @@ def constant_fold(
         attempting constant folding, and any old non-shared inputs will not work with
         the returned outputs
     """
-    fg = FunctionGraph(outputs=xs, features=[ShapeFeature()], clone=True)
+    fg = FunctionGraph(outputs=xs, features=[ShapeFeature()], copy_inputs=False, clone=True)
 
     # By default, rewrite_graph includes canonicalize which includes constant-folding as the final rewrite
     folded_xs = rewrite_graph(fg).outputs


Test patch:
diff --git a/tests/model/transform/test_optimization.py b/tests/model/transform/test_optimization.py
index 9866c1b62..9b697f630 100644
--- a/tests/model/transform/test_optimization.py
+++ b/tests/model/transform/test_optimization.py
@@ -20,8 +20,10 @@ from pytensor.graph import Constant
 from pymc import Deterministic, do
 from pymc.data import Data
 from pymc.distributions import HalfNormal, Normal
+from pymc.exceptions import NotConstantValueError
 from pymc.model import Model
 from pymc.model.transform.optimization import freeze_dims_and_data
+from pymc.pytensorf import constant_fold
 
 
 def test_freeze_dims_and_data():
@@ -144,3 +146,16 @@ def test_freeze_dim_after_do_intervention():
 
     frozen_do_m = freeze_dims_and_data(do_m)
     assert frozen_do_m["x"].type.shape == (5,)
+
+
+def test_freeze_dims_and_data_partially_observed_rv():
+    # Regression test for #7387
+
+    with Model(coords={"a": [0, 1, 2]}) as model:
+        y = Normal("y", 0, observed=[0, 0, np.nan], dims="a")
+
+    with pytest.raises(NotConstantValueError):
+        constant_fold([y.shape])
+
+    frozen_y = freeze_dims_and_data(model)["y"]
+    assert constant_fold([frozen_y.shape]) == (3,)
diff --git a/tests/test_pytensorf.py b/tests/test_pytensorf.py
index f6084718f..190b159db 100644
--- a/tests/test_pytensorf.py
+++ b/tests/test_pytensorf.py
@@ -646,25 +646,33 @@ def test_reseed_rngs():
         assert rng.get_value().bit_generator.state == bit_generator.state
 
 
-def test_constant_fold():
-    x = pt.random.normal(size=(5,))
-    y = pt.arange(x.size)
-
-    res = constant_fold((y, y.shape))
-    assert np.array_equal(res[0], np.arange(5))
-    assert tuple(res[1]) == (5,)
-
-
-def test_constant_fold_raises():
-    size = pytensor.shared(5)
-    x = pt.random.normal(size=(size,))
-    y = pt.arange(x.size)
-
-    with pytest.raises(NotConstantValueError):
-        constant_fold((y, y.shape))
-
-    res = constant_fold((y, y.shape), raise_not_constant=False)
-    assert tuple(res[1].eval()) == (5,)
+class TestConstantFold:
+    def test_constant_fold(self):
+        x = pt.random.normal(size=(5,))
+        y = pt.arange(x.size)
+
+        res = constant_fold((y, y.shape))
+        assert np.array_equal(res[0], np.arange(5))
+        assert tuple(res[1]) == (5,)
+
+    def test_constant_fold_raises(self):
+        size = pytensor.shared(5)
+        x = pt.random.normal(size=(size,))
+        y = pt.arange(x.size)
+
+        with pytest.raises(NotConstantValueError):
+            constant_fold((y, y.shape))
+
+        res = constant_fold((y, y.shape), raise_not_constant=False)
+        assert tuple(res[1].eval()) == (5,)
+
+    def test_inputs_preserved(self):
+        # Make sure constant_folded graph depends on original graph inputs (not copies)
+        # Regression test for #7387
+        a = pt.scalar("a", dtype="int")
+        out = pt.empty((a,))
+        (out_shape,) = constant_fold((out.shape[0],), raise_not_constant=False)
+        assert out_shape is a
 
 
 def test_replace_vars_in_graphs():