Fix shape_inference of ChoiceRV when param_shapes are provided

Author: ricardoV94

pytensor/tensor/random/basic.py

@@ -1990,11 +1990,12 @@ def _supp_shape_from_params(self, *args, **kwargs):
         raise NotImplementedError()
 
     def _infer_shape(self, size, dist_params, param_shapes=None):
-        (a, p, _) = dist_params
-
+        a, p, _ = dist_params
         if isinstance(p.type, pytensor.tensor.type_other.NoneTypeT):
+            param_shapes = param_shapes[:1] if param_shapes is not None else None
             shape = super()._infer_shape(size, (a,), param_shapes)
         else:
+            param_shapes = param_shapes[:2] if param_shapes is not None else None
             shape = super()._infer_shape(size, (a, p), param_shapes)
 
         return shape

tests/tensor/random/test_basic.py

@@ -1390,6 +1390,19 @@ def test_choice_samples():
     compare_sample_values(choice, at.as_tensor_variable([1, 2, 3]), 2, replace=True)
 
 
+def test_choice_infer_shape():
+    node = choice([0, 1]).owner
+    res = node.op._infer_shape((), node.inputs[3:], None)
+    assert tuple(res.eval()) == ()
+
+    node = choice([0, 1]).owner
+    # The param_shape of a NoneConst is None, during shape_inference
+    res = node.op._infer_shape(
+        (), node.inputs[3:], (node.inputs[3].shape, None, node.inputs[5].shape)
+    )
+    assert tuple(res.eval()) == ()
+
+
 def test_permutation_samples():
     compare_sample_values(
         permutation,