.wip

Author: ricardoV94

tests/xtensor/test_random.py

@@ -36,7 +36,156 @@ def test_all_basic_rvs_are_wrapped():
 
 
 def test_normal():
-    pass
+    c_size = tensor("c_size", shape=(), dtype=int)
+    c_size_xr = as_xtensor(c_size, name="c_size_xr")
+
+    # Vector inputs
+    mu_vec_in = tensor("mu", shape=(2,))
+    sigma_vec_in = tensor("sigma", shape=(2,))
+
+    mu_vec_xr = as_xtensor(mu_vec_in, dims=("a",), name="mu_xr")
+    sigma_vec_xr = as_xtensor(sigma_vec_in, dims=("a",), name="sigma_xr")
+
+    # Vector inputs: Basic case (no extra_dims)
+    out_vec = pxr.normal(mu_vec_xr, sigma_vec_xr, rng=rng)
+    assert out_vec.type.dims == ("a",)
+    assert out_vec.type.shape == (2,)
+    assert equal_computations(
+        [lower_rewrite(out_vec.values)],
+        [rewrite_graph(ptr.normal(mu_vec_in, sigma_vec_in, rng=rng))],
+    )
+
+    mu_val = np.array([0.0, 10.0])
+    sigma_val = np.array([1.0, 2.0])
+    eval_rng_seed_vec_basic = 12345
+
+    actual_val_vec_basic = out_vec.eval(
+        {
+            mu_vec_in: mu_val,
+            sigma_vec_in: sigma_val,
+            rng: np.random.default_rng(eval_rng_seed_vec_basic),
+        }
+    )
+    expected_val_vec_basic = np.random.default_rng(eval_rng_seed_vec_basic).normal(
+        mu_val, sigma_val
+    )
+    np.testing.assert_allclose(actual_val_vec_basic, expected_val_vec_basic)
+
+    # Vector inputs: With extra_dims
+    out_vec_extra = pxr.normal(
+        mu_vec_xr, sigma_vec_xr, extra_dims=dict(c=c_size_xr), rng=rng
+    )
+    assert out_vec_extra.type.dims == ("c", "a")
+    assert equal_computations(
+        [lower_rewrite(out_vec_extra.values)],
+        [
+            rewrite_graph(
+                ptr.normal(
+                    mu_vec_in, sigma_vec_in, size=(c_size, mu_vec_in.shape[0]), rng=rng
+                )
+            )
+        ],
+    )
+
+    c_size_val = 5
+    eval_rng_seed_vec_extra = 67890
+    actual_val_vec_extra = out_vec_extra.eval(
+        {
+            mu_vec_in: mu_val,
+            sigma_vec_in: sigma_val,
+            c_size: c_size_val,
+            rng: np.random.default_rng(eval_rng_seed_vec_extra),
+        }
+    )
+    expected_val_vec_extra = np.random.default_rng(eval_rng_seed_vec_extra).normal(
+        loc=mu_val, scale=sigma_val, size=(c_size_val, mu_val.shape[0])
+    )
+    np.testing.assert_allclose(actual_val_vec_extra, expected_val_vec_extra)
+
+    # Scalar inputs
+    mu_scalar_in = tensor("mu_s", shape=())
+    sigma_scalar_in = tensor("sigma_s", shape=())
+
+    mu_scalar_xr = as_xtensor(mu_scalar_in, name="mu_s_xr")
+    sigma_scalar_xr = as_xtensor(sigma_scalar_in, name="sigma_s_xr")
+
+    # Scalar inputs: Basic case
+    out_scalar = pxr.normal(mu_scalar_xr, sigma_scalar_xr, rng=rng)
+    assert out_scalar.type.dims == ()
+    assert out_scalar.type.shape == ()
+    assert equal_computations(
+        [lower_rewrite(out_scalar.values)],
+        [rewrite_graph(ptr.normal(mu_scalar_in, sigma_scalar_in, rng=rng))],
+    )
+
+    mu_s_val = 0.0
+    sigma_s_val = 1.0
+    eval_rng_seed_scalar_basic = 23456
+    actual_val_scalar_basic = out_scalar.eval(
+        {
+            mu_scalar_in: mu_s_val,
+            sigma_scalar_in: sigma_s_val,
+            rng: np.random.default_rng(eval_rng_seed_scalar_basic),
+        }
+    )
+    expected_val_scalar_basic = np.random.default_rng(
+        eval_rng_seed_scalar_basic
+    ).normal(mu_s_val, sigma_s_val)
+    np.testing.assert_allclose(actual_val_scalar_basic, expected_val_scalar_basic)
+
+    # Scalar inputs: With extra_dims
+    out_scalar_extra = pxr.normal(
+        mu_scalar_xr, sigma_scalar_xr, extra_dims=dict(c=c_size_xr), rng=rng
+    )
+    assert out_scalar_extra.type.dims == ("c",)
+    assert equal_computations(
+        [lower_rewrite(out_scalar_extra.values)],
+        [
+            rewrite_graph(
+                ptr.normal(mu_scalar_in, sigma_scalar_in, size=(c_size,), rng=rng)
+            )
+        ],
+    )
+
+    eval_rng_seed_scalar_extra = 78901
+    actual_val_scalar_extra = out_scalar_extra.eval(
+        {
+            mu_scalar_in: mu_s_val,
+            sigma_scalar_in: sigma_s_val,
+            c_size: c_size_val,
+            rng: np.random.default_rng(eval_rng_seed_scalar_extra),
+        }
+    )
+    expected_val_scalar_extra = np.random.default_rng(
+        eval_rng_seed_scalar_extra
+    ).normal(loc=mu_s_val, scale=sigma_s_val, size=(c_size_val,))
+    np.testing.assert_allclose(actual_val_scalar_extra, expected_val_scalar_extra)
+
+    # Error conditions
+    # Invalid core_dims: normal is element-wise, expects core_dims=() for params.
+    with pytest.raises(
+        ValueError,
+        match=re.escape(
+            "Parameter mu_xr has invalid core dimensions ['a']. "
+            "Expected [] based on RV definition and core_dims argument."
+        ),
+    ):
+        pxr.normal(mu_vec_xr, sigma_vec_xr, core_dims=("a",), rng=rng)
+
+    # Invalid extra_dims (conflicting with existing batch dims)
+    a_size_xr = mu_vec_xr.sizes["a"]
+    with pytest.raises(
+        ValueError,
+        match=re.escape(
+            "Size dimensions ['a'] conflict with parameter dimensions. They should be unique."
+        ),
+    ):
+        pxr.normal(
+            mu_vec_xr,
+            sigma_vec_xr,
+            extra_dims=dict(c=c_size_xr, a=a_size_xr),  # 'a' conflicts
+            rng=rng,
+        )
 
 
 def test_categorical():