Preserve initialized loss dependencies and avoid cached JAX tracers

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'])

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,41 @@ 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.]))
+
+
+@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))