.WIP

Author: ricardoV94

pytensor/xtensor/vectorization.py

@@ -117,3 +117,23 @@ def make_node(self, *inputs):
             for core_out, core_out_dims in zip(core_node.outputs, core_outputs_dims)
         ]
         return Apply(self, inputs, outputs)
+
+
+class XRandomVariable(XOp):
+    __props__ = ("dist", "core_dims")
+
+    def __init__(self, dist, core_dims: tuple[tuple[tuple[str, ...], ...], tuple[tuple[str, ...], ...]]):
+        super().__init__()
+        self.dist = dist
+        self.core_dims = core_dims,
+
+    def make_node(self, *inputs):
+        inputs = [as_xtensor(inp) for inp in inputs]
+        if len(inputs) != self.dist.nin:
+            raise ValueError(
+                f"Wrong number of inputs, expected {self.dist.nin}, got {len(inputs)}"
+            )
+
+        output_dims, output_shape = inputs[0].type.dims, inputs[0].type.shape
+        output = xtensor(dtype=self.dist.dtype, dims=output_dims, shape=output_shape)
+        return Apply(self, inputs, [output])

tests/xtensor/test_random.py

@@ -0,0 +1,73 @@
+import pytest
+
+import pytensor.tensor.random as ptr
+from pytensor.graph.basic import equal_computations
+from pytensor.tensor.random.type import random_generator_type
+from pytensor.xtensor import xtensor
+from pytensor.xtensor.random import multinomial, multivariate_normal, normal, categorical
+
+lower_rewrite = lambda x: x
+
+def test_normal():
+    pass
+
+def test_categorical():
+    pass
+
+def test_multinomial():
+    rng = random_generator_type("rng")
+    n = xtensor(shape=(2,), dims=("a",))
+    p = xtensor(shape=(3, None), dims=("p", "a"))
+    c_size = xtensor(shape=(), dims=(), dtype=int)
+    a_size = n.sizes["a"]
+
+    out = multinomial(n, p, core_dims=("p",), rng=rng)
+    assert out.type.dims == ("a", "p")
+    assert out.type.shape == (2, 3)
+    assert equal_computations(
+        [lower_rewrite(out)],
+        [ptr.multinomial(n.values, p.values.T, size=None, rng=rng)],
+    )
+    # TODO: Make sure we can actually evaluate it
+    ...
+
+    out = multinomial(n, p, core_dims=("p",), size=dict(a=a_size), rng=rng)
+    assert out.type.dims == ("a", "p")
+    assert equal_computations(
+        [lower_rewrite(out)],
+        [ptr.multinomial(n.values, p.values.T, size=(a_size.values,), rng=rng)],
+    )
+
+    out = multinomial(n, p, core_dims=("p",), size=dict(a=a_size, c=c_size), rng=rng)
+    assert out.type.dims == ("a", "c", "p")
+    assert equal_computations(
+        [lower_rewrite(out)],
+        [ptr.multinomial(n.values[:, None], p.values.T[:, None, :], size=(a_size.values, c_size.values), rng=rng)],
+    )
+
+    out = multinomial(n, p, core_dims=("p",), size=dict(c=c_size, a=a_size,), rng=rng)
+    assert out.type.dims == ("c", "a", "p")
+    assert equal_computations(
+        [lower_rewrite(out)],
+        [ptr.multinomial(n.values, p.values.T, size=(c_size.values, a_size.values), rng=rng)],
+    )
+
+    # Test missing core_dims
+    with pytest.raises(ValueError):
+        multinomial(n, p, rng=rng)
+
+    # Test invalid core_dims
+    with pytest.raises(ValueError):
+        # n cannot have a core dimension
+        multinomial(n, p, core_dims=("a",), rng=rng)
+
+    # Test incomplete size
+    with pytest.raises(ValueError):
+        multinomial(n, p, core_dims=("p",), size=dict(c=c_size), rng=rng)
+
+
+def test_multivariate_normal():
+    pass
+
+def test_new_out_dim()
+    pass