Make rv_size_is_none more robust

Author: ricardoV94

pymc/distributions/shape_utils.py

@@ -22,7 +22,7 @@
 
 import numpy as np
 
-from aesara.graph.basic import Constant, Variable
+from aesara.graph.basic import Variable
 from aesara.tensor.var import TensorVariable
 from typing_extensions import TypeAlias
 
@@ -618,4 +618,4 @@ def find_size(
 
 def rv_size_is_none(size: Variable) -> bool:
     """Check wether an rv size is None (ie., at.Constant([]))"""
-    return isinstance(size, Constant) and size.data.size == 0
+    return size.type.shape == (0,)

pymc/tests/test_distributions_moments.py

@@ -133,6 +133,9 @@ def test_rv_size_is_none():
     rv = Normal.dist(0, 1, size=None)
     assert rv_size_is_none(rv.owner.inputs[1])
 
+    rv = Normal.dist(0, 1, size=())
+    assert rv_size_is_none(rv.owner.inputs[1])
+
     rv = Normal.dist(0, 1, size=1)
     assert not rv_size_is_none(rv.owner.inputs[1])