Add test for compound binning correction

tests/test_base.py

@@ -183,18 +183,6 @@ def test_unsupported_correction():
         CorrectionWithGradient(schemas["categorical"])
 
 
-def test_unsupported_compound_binning():
-    with pytest.raises(
-        ValueError,
-        match=(
-            "Correction 'compound non-uniform binning' contains a compound "
-            "Binning correction \\(one or more of the bin contents are not "
-            "simple scalars\\). This is not supported."
-        ),
-    ):
-        CorrectionWithGradient(schemas["compound-nonuniform-binning"])
-
-
 def test_unsupported_flow_type():
     with pytest.raises(
         ValueError,
@@ -428,3 +416,13 @@ def test_formula_with_parameters(jit):
     assert np.allclose(values, [5.0, 7.0])
     assert len(grads) == 2
     assert np.allclose(grads, [2.0, 2.0])
+
+
+def test_unsupported_compound_binning():
+    cg = CorrectionWithGradient(schemas["compound-nonuniform-binning"])
+    value = cg.evaluate(0.5)
+    assert math.isclose(value, 0.5 * 0.5)
+
+    value, grad = jax.value_and_grad(cg.evaluate)(0.5)
+    assert math.isclose(value, 0.5 * 0.5)
+    assert math.isclose(grad, 1.0)