more versions of relative entropy

Author: dance858

sparsediffpy/_core/_nodes_bivariate.py

@@ -32,10 +32,24 @@ def __init__(self, x, z):
 
 
 class RelEntr(Expression):
-    """x * log(x / y) elementwise."""
+    """x * log(x / y) elementwise.
+
+    Supports three variants (auto-dispatched by the C layer):
+    - Both same shape: elementwise
+    - Scalar x, vector y: x * log(x / y_i) for each i
+    - Vector x, scalar y: x_i * log(x_i / y) for each i
+    """
     def __init__(self, x, y):
-        if x.shape != y.shape:
-            raise ValueError(f"rel_entr: shape mismatch {x.shape} vs {y.shape}")
+        if x.shape == y.shape:
+            self.shape = x.shape
+        elif is_scalar(x.shape):
+            self.shape = y.shape
+        elif is_scalar(y.shape):
+            self.shape = x.shape
+        else:
+            raise ValueError(
+                f"rel_entr: shapes must match or one must be scalar, "
+                f"got {x.shape} and {y.shape}"
+            )
         self.x = x
         self.y = y
-        self.shape = x.shape

tests/other/test_rel_entr.py

@@ -1,9 +1,12 @@
 import numpy as np
+import pytest
 import sparsediffpy as sp
 from tests.utils import NumericalDerivativeChecker, random_positive_point
 
 
-def test_rel_entr_jacobian(scope, rng):
+# --- Vector-vector (both same shape) ---
+
+def test_rel_entr_vector_jacobian(scope, rng):
     x = scope.Variable(3, 1)
     y = scope.Variable(3, 1)
     f = sp.rel_entr(x, y)
@@ -12,7 +15,7 @@ def test_rel_entr_jacobian(scope, rng):
     checker.check_jacobian(random_positive_point(scope, rng))
 
 
-def test_rel_entr_hessian(scope, rng):
+def test_rel_entr_vector_hessian(scope, rng):
     x = scope.Variable(3, 1)
     y = scope.Variable(3, 1)
     f = sp.rel_entr(x, y)
@@ -22,7 +25,7 @@ def test_rel_entr_hessian(scope, rng):
     checker.check_hessian(x0, rng.standard_normal(3))
 
 
-def test_rel_entr_forward(scope, rng):
+def test_rel_entr_vector_forward(scope, rng):
     x = scope.Variable(3, 1)
     y = scope.Variable(3, 1)
     f = sp.rel_entr(x, y)
@@ -32,3 +35,80 @@ def test_rel_entr_forward(scope, rng):
     y_val = y.value
     expected = x_val * np.log(x_val / y_val)
     np.testing.assert_allclose(fn.forward(), expected)
+
+
+# --- Scalar x, vector y ---
+
+def test_rel_entr_scalar_vector_jacobian(scope, rng):
+    x = scope.Variable(1, 1)
+    y = scope.Variable(3, 1)
+    f = sp.rel_entr(x, y)
+    assert f.shape == (3, 1)
+    fn = sp.compile(f)
+    checker = NumericalDerivativeChecker(fn, scope)
+    checker.check_jacobian(random_positive_point(scope, rng))
+
+
+def test_rel_entr_scalar_vector_hessian(scope, rng):
+    x = scope.Variable(1, 1)
+    y = scope.Variable(3, 1)
+    f = sp.rel_entr(x, y)
+    fn = sp.compile(f)
+    checker = NumericalDerivativeChecker(fn, scope)
+    x0 = random_positive_point(scope, rng)
+    checker.check_hessian(x0, rng.standard_normal(3))
+
+
+def test_rel_entr_scalar_vector_forward(scope, rng):
+    x = scope.Variable(1, 1)
+    y = scope.Variable(3, 1)
+    f = sp.rel_entr(x, y)
+    fn = sp.compile(f)
+    x0 = random_positive_point(scope, rng)
+    x_val = x.value[0]
+    y_val = y.value
+    expected = x_val * np.log(x_val / y_val)
+    np.testing.assert_allclose(fn.forward(), expected)
+
+
+# --- Vector x, scalar y ---
+
+def test_rel_entr_vector_scalar_jacobian(scope, rng):
+    x = scope.Variable(3, 1)
+    y = scope.Variable(1, 1)
+    f = sp.rel_entr(x, y)
+    assert f.shape == (3, 1)
+    fn = sp.compile(f)
+    checker = NumericalDerivativeChecker(fn, scope)
+    checker.check_jacobian(random_positive_point(scope, rng))
+
+
+def test_rel_entr_vector_scalar_hessian(scope, rng):
+    x = scope.Variable(3, 1)
+    y = scope.Variable(1, 1)
+    f = sp.rel_entr(x, y)
+    fn = sp.compile(f)
+    checker = NumericalDerivativeChecker(fn, scope)
+    x0 = random_positive_point(scope, rng)
+    checker.check_hessian(x0, rng.standard_normal(3))
+
+
+def test_rel_entr_vector_scalar_forward(scope, rng):
+    x = scope.Variable(3, 1)
+    y = scope.Variable(1, 1)
+    f = sp.rel_entr(x, y)
+    fn = sp.compile(f)
+    x0 = random_positive_point(scope, rng)
+    x_val = x.value
+    y_val = y.value[0]
+    expected = x_val * np.log(x_val / y_val)
+    np.testing.assert_allclose(fn.forward(), expected)
+
+
+# --- Shape mismatch ---
+
+def test_rel_entr_incompatible_shapes(scope):
+    x = scope.Variable(3, 1)
+    y = scope.Variable(2, 1)
+    with pytest.raises(ValueError, match="shapes must match or one must be scalar"):
+        sp.rel_entr(x, y)

tests/test_validation.py

@@ -44,7 +44,7 @@ def test_matmul_inner_dim_mismatch(self, scope):
     def test_rel_entr_shape_mismatch(self, scope):
         x = scope.Variable(3, 1)
         y = scope.Variable(2, 1)
-        with pytest.raises(ValueError, match="shape mismatch"):
+        with pytest.raises(ValueError, match="shapes must match or one must be scalar"):
             sp.rel_entr(x, y)
 
     def test_quad_over_lin_non_scalar_z(self, scope):