diff --git a/spacecore/functional/_base.py b/spacecore/functional/_base.py index 8afa2da..f4e0f29 100644 --- a/spacecore/functional/_base.py +++ b/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": """ diff --git a/tests/functional/test_functional_base.py b/tests/functional/test_functional_base.py index df075b7..6678745 100644 --- a/tests/functional/test_functional_base.py +++ b/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 # ===========================================================================