From 2f3a353a5042d5eedc28c62acfd3637a556cd7fa Mon Sep 17 00:00:00 2001 From: Rogerio Jorge Date: Fri, 23 Oct 2026 17:57:12 -0500 Subject: [PATCH 1/2] Preserve loss dependencies and avoid cached JAX tracers --- essos/losses.py | 19 +++++++++++++------ tests/test_multiobjectives.py | 17 +++++++++++++++++ 2 files changed, 30 insertions(+), 6 deletions(-) diff --git a/essos/losses.py b/essos/losses.py index 16360d13..949529d6 100644 --- a/essos/losses.py +++ b/essos/losses.py @@ -1,6 +1,6 @@ from functools import partial import jax.numpy as jnp -from jax import tree_util, jit, grad as jax_grad, value_and_grad as jax_value_and_grad +from jax import tree_util, jit, ensure_compile_time_eval, grad as jax_grad, value_and_grad as jax_value_and_grad from jax.flatten_util import ravel_pytree class base_loss: @@ -79,9 +79,13 @@ def clear_cache(self): def _ensure_unravelers(self): if self._starting_dofs is None or self._dofs_to_args is None or self._dofs_to_pytree is None: - self._starting_dofs, tuple_unraveler = ravel_pytree( - tuple(self.dependencies[arg] for arg in self.args_names) - ) + # Dependencies are constants of this loss. If the first call is + # jitted, build the cached vector eagerly rather than retaining a + # tracer that would escape the compiled function. + with ensure_compile_time_eval(): + self._starting_dofs, tuple_unraveler = ravel_pytree( + tuple(self.dependencies[arg] for arg in self.args_names) + ) self._dofs_to_args = tuple_unraveler def _named_unraveler(dofs): @@ -153,6 +157,7 @@ def __mul__(self, other): new_fun = lambda *args, **kwargs: other * self.fun(*args, **kwargs) out_loss = custom_loss(new_fun, *self.args_names, **self.kwargs) + out_loss.dependencies = self.dependencies return out_loss @@ -185,13 +190,15 @@ def dependencies(self, value): @property def starting_dofs(self): if self._starting_dofs is None: - self._starting_dofs, self._dofs_to_pytree = ravel_pytree(self.dependencies) + with ensure_compile_time_eval(): + self._starting_dofs, self._dofs_to_pytree = ravel_pytree(self.dependencies) return self._starting_dofs @property def dofs_to_pytree(self): if self._dofs_to_pytree is None: - self._starting_dofs, self._dofs_to_pytree = ravel_pytree(self.dependencies) + with ensure_compile_time_eval(): + self._starting_dofs, self._dofs_to_pytree = ravel_pytree(self.dependencies) return self._dofs_to_pytree @partial(jit, static_argnames=['self']) diff --git a/tests/test_multiobjectives.py b/tests/test_multiobjectives.py index 7c53cbff..c84f4d5b 100644 --- a/tests/test_multiobjectives.py +++ b/tests/test_multiobjectives.py @@ -1,4 +1,5 @@ import pytest +import numpy as np from unittest.mock import MagicMock, patch import essos.objective_functions as objf from essos.losses import custom_loss @@ -151,3 +152,19 @@ def test_plots_without_plotly_raise_clear_error(monkeypatch): monkeypatch.setitem(sys.modules, "plotly", None) with pytest.raises(ImportError, match="essos\\[multiobjective\\]"): MultiObjectiveOptimizer.plot_param_importances(None) + + + +def test_scaling_an_initialized_custom_loss_preserves_dependencies(): + loss = custom_loss(lambda x: jnp.sum(x*x), 'x') + loss.dependencies = {'x': jnp.array([2., 3.])} + scaled = 3*loss + np.testing.assert_allclose(scaled(scaled.starting_dofs), 39.) + np.testing.assert_allclose(scaled.grad(scaled.starting_dofs), jnp.array([12., 18.])) + + +def test_first_jitted_call_does_not_cache_traced_starting_dofs(): + loss = custom_loss(lambda x, y: jnp.sum(x*x)+jnp.sum(y), 'x', 'y') + loss.dependencies = {'x': jnp.array([2., 3.]), 'y': jnp.array([4.])} + np.testing.assert_allclose(loss(jnp.array([2., 3., 4.])), 17.) + np.testing.assert_allclose(loss.starting_dofs + 0., jnp.array([2., 3., 4.])) From fd2bb49bfb84ade3514688616d132107e90ea99b Mon Sep 17 00:00:00 2001 From: Rogerio Jorge Date: Fri, 23 Oct 2026 18:29:05 -0500 Subject: [PATCH 2/2] Cover composite initialization and shared dependency gradients --- tests/test_multiobjectives.py | 22 ++++++++++++++++++++++ 1 file changed, 22 insertions(+) diff --git a/tests/test_multiobjectives.py b/tests/test_multiobjectives.py index c84f4d5b..e2ca3691 100644 --- a/tests/test_multiobjectives.py +++ b/tests/test_multiobjectives.py @@ -168,3 +168,25 @@ def test_first_jitted_call_does_not_cache_traced_starting_dofs(): loss.dependencies = {'x': jnp.array([2., 3.]), 'y': jnp.array([4.])} np.testing.assert_allclose(loss(jnp.array([2., 3., 4.])), 17.) np.testing.assert_allclose(loss.starting_dofs + 0., jnp.array([2., 3., 4.])) + + +@pytest.mark.parametrize('first_access', ['call', 'starting_dofs']) +def test_composite_first_jit_access_preserves_dofs_and_dependency_gradients(first_access): + import jax + + quadratic = custom_loss(lambda x, y: jnp.sum(x*x) + jnp.sum(y), 'x', 'y') + coupling = custom_loss(lambda x: x[0]*x[1], 'x') + total = quadratic + coupling + total.dependencies = {'x': jnp.array([2., 3.]), 'y': jnp.array([4.]), + 'unused': jnp.array([99.])} + # Dictionary pytrees flatten their keys in sorted order. + dofs = jnp.array([99., 2., 3., 4.]) + if first_access == 'call': + np.testing.assert_allclose(total(dofs), 23.) + else: + np.testing.assert_allclose(jax.jit(lambda: total.starting_dofs)(), dofs) + # The cached vector must remain usable outside the initial JIT trace. + np.testing.assert_allclose(total.starting_dofs + 0., dofs) + np.testing.assert_allclose(total(dofs), 23.) + np.testing.assert_allclose(total.grad(dofs), jnp.array([0., 7., 8., 1.])) + np.testing.assert_allclose(jax.grad(total)(dofs), total.grad(dofs))