functional: add value_and_grad and auxiliary value/grad args

spacecore/functional/_base.py

@@ -58,17 +58,33 @@ def domain(self) -> Domain:
         return self.dom
 
     @abstractmethod
-    def value(self, x: Any) -> Any:
-        """Evaluate this functional at an element of ``self.domain``."""
+    def value(self, x: Any, *args: Any, **kwargs: Any) -> Any:
+        """Evaluate this functional at an element of ``self.domain``.
 
-    def grad(self, x: Any) -> Any:
+        Subclasses may accept extra positional/keyword arguments — auxiliary
+        parameters such as data, temperature, or a penalty weight — that are not
+        part of the domain. Overrides that take only ``x`` remain valid.
+        """
+
+    def grad(self, x: Any, *args: Any, **kwargs: Any) -> Any:
         """Gradient at an element of ``self.domain``.
 
         Override in subclasses that support differentiation; the base raises
-        :class:`NotImplementedError`.
+        :class:`NotImplementedError`. Any auxiliary ``*args``/``**kwargs`` mirror
+        those accepted by :meth:`value`.
         """
         raise NotImplementedError(f"{type(self).__name__} does not implement grad().")
 
+    def value_and_grad(self, x: Any, *args: Any, **kwargs: Any) -> tuple[Any, Any]:
+        """Return ``(value, gradient)`` at ``x``.
+
+        The base default evaluates :meth:`value` and :meth:`grad` separately.
+        Subclasses may override with a single-pass evaluator (e.g.
+        ``jax.value_and_grad``); an override must return the same pair as the
+        default, with the gradient in the same geometry as :meth:`grad`.
+        """
+        return self.value(x, *args, **kwargs), self.grad(x, *args, **kwargs)
+
     def vgrad(self, xs: Any) -> Any:
         """Gradient over a leading batch axis.
 
@@ -77,13 +93,13 @@ def vgrad(self, xs: Any) -> Any:
         """
         raise NotImplementedError(f"{type(self).__name__} does not implement vgrad().")
 
-    def _value_core(self, x: Any) -> Any:
+    def _value_core(self, x: Any, *args: Any, **kwargs: Any) -> Any:
         """Check-free value core; the base falls back to the checked ``value``."""
-        return self.value(x)
+        return self.value(x, *args, **kwargs)
 
-    def _grad_core(self, x: Any) -> Any:
+    def _grad_core(self, x: Any, *args: Any, **kwargs: Any) -> Any:
         """Check-free gradient core; the base falls back to the checked ``grad``."""
-        return self.grad(x)
+        return self.grad(x, *args, **kwargs)
 
     def _vvalue_core(self, xs: Any) -> Any:
         """Check-free batched-value core; the base falls back to ``vvalue``."""
@@ -93,9 +109,9 @@ def _vgrad_core(self, xs: Any) -> Any:
         """Check-free batched-gradient core; the base falls back to ``vgrad``."""
         return self.vgrad(xs)
 
-    def __call__(self, x: Any) -> Any:
+    def __call__(self, x: Any, *args: Any, **kwargs: Any) -> Any:
         """Evaluate this functional at ``x``."""
-        return self.value(x)
+        return self.value(x, *args, **kwargs)
 
     def compose(self, A: "LinOp") -> "Functional":
         """

tests/functional/test_functional_base.py

@@ -52,6 +52,27 @@ def _convert(self, new_ctx):
         return _SumSquares(self.domain.convert(new_ctx), new_ctx)
 
 
+class _ScaledSumSquares(_SumSquares):
+    """``F(x, scale=1.0) = scale * sum(x * x)`` — exercises auxiliary value args."""
+
+    def value(self, x, scale=1.0):
+        return self.ops.sum(x * x) * scale
+
+
+class _SumSquaresGrad(_SumSquares):
+    """``_SumSquares`` with a gradient, so ``value_and_grad`` has both halves."""
+
+    def grad(self, x):
+        return 2.0 * x
+
+
+class _SumSquaresFused(_SumSquaresGrad):
+    """Overrides ``value_and_grad`` with a single-pass evaluator."""
+
+    def value_and_grad(self, x):
+        return self.ops.sum(x * x), 2.0 * x
+
+
 # ===========================================================================
 # Abstract enforcement
 # ===========================================================================
@@ -82,6 +103,54 @@ def test_call_is_alias_for_value(self, numpy_ctx):
         np.testing.assert_allclose(to_numpy(f.value(x)), 14.0)
 
 
+# ===========================================================================
+# W0/F1: value(x, *args, **kwargs) — auxiliary parameters forwarded
+# ===========================================================================
+class TestValueAuxArgs:
+    def test_value_accepts_positional_and_keyword(self, numpy_ctx):
+        space = sc.DenseCoordinateSpace((3,), numpy_ctx)
+        f = _ScaledSumSquares(space, numpy_ctx)
+        x = numpy_ctx.asarray([1.0, 2.0, 3.0])
+        np.testing.assert_allclose(to_numpy(f.value(x)), 14.0)  # default scale
+        np.testing.assert_allclose(to_numpy(f.value(x, 2.0)), 28.0)  # positional
+        np.testing.assert_allclose(to_numpy(f.value(x, scale=0.5)), 7.0)  # keyword
+
+    def test_call_forwards_aux_args(self, numpy_ctx):
+        space = sc.DenseCoordinateSpace((3,), numpy_ctx)
+        f = _ScaledSumSquares(space, numpy_ctx)
+        x = numpy_ctx.asarray([1.0, 2.0, 3.0])
+        np.testing.assert_allclose(to_numpy(f(x, scale=3.0)), 42.0)
+        np.testing.assert_allclose(to_numpy(f(x, 3.0)), 42.0)
+
+
+# ===========================================================================
+# W0/F2: value_and_grad
+# ===========================================================================
+class TestValueAndGrad:
+    def test_default_equals_value_then_grad(self, numpy_ctx):
+        space = sc.DenseCoordinateSpace((3,), numpy_ctx)
+        f = _SumSquaresGrad(space, numpy_ctx)
+        x = numpy_ctx.asarray([1.0, 2.0, 3.0])
+        value, grad = f.value_and_grad(x)
+        np.testing.assert_allclose(to_numpy(value), to_numpy(f.value(x)))
+        np.testing.assert_allclose(to_numpy(grad), to_numpy(f.grad(x)))
+
+    def test_fused_override_matches_default(self, numpy_ctx):
+        space = sc.DenseCoordinateSpace((3,), numpy_ctx)
+        x = numpy_ctx.asarray([1.0, 2.0, 3.0])
+        default = _SumSquaresGrad(space, numpy_ctx).value_and_grad(x)
+        fused = _SumSquaresFused(space, numpy_ctx).value_and_grad(x)
+        np.testing.assert_allclose(to_numpy(fused[0]), to_numpy(default[0]))
+        np.testing.assert_allclose(to_numpy(fused[1]), to_numpy(default[1]))
+
+    def test_default_without_grad_raises(self, numpy_ctx):
+        space = sc.DenseCoordinateSpace((3,), numpy_ctx)
+        f = _SumSquares(space, numpy_ctx)  # no grad defined
+        x = numpy_ctx.asarray([1.0, 2.0, 3.0])
+        with pytest.raises(NotImplementedError):
+            f.value_and_grad(x)
+
+
 # ===========================================================================
 # Explicit context priority
 # ===========================================================================