Add logic for networks with "is_training" (e.g. don't want dropout when predicting)

Author: asmith26

jax_toolkit/losses/tests/test_utils.py

@@ -28,9 +28,9 @@ def net_function(x: jnp.ndarray) -> jnp.ndarray:
         rng = jax.random.PRNGKey(42)
         params = net_transform.init(rng, jnp.array(0))
 
-        self.assertEqual(0, actual_loss_function_wrapper(params, x=jnp.array(0), y_true=jnp.array(0)))
-        self.assertEqual(0, actual_loss_function_wrapper(params, x=jnp.array(1), y_true=jnp.array(1)))
-        self.assertEqual(1, actual_loss_function_wrapper(params, x=jnp.array(0), y_true=jnp.array(1)))
+        self.assertEqual(0, actual_loss_function_wrapper(params, jnp.array(0), jnp.array(0)))
+        self.assertEqual(0, actual_loss_function_wrapper(params, jnp.array(1), jnp.array(1)))
+        self.assertEqual(1, actual_loss_function_wrapper(params, jnp.array(0), jnp.array(1)))
 
     def test_supported_loss_returns_correctly_with_loss_kwargs(self):
         import haiku as hk

jax_toolkit/losses/utils.py

@@ -1,3 +1,4 @@
+from functools import partial
 from typing import Callable, Dict, Optional
 
 import jax
@@ -43,16 +44,19 @@
 
 def get_haiku_loss_function(
     net_transform: hk.Transformed, loss: str, **loss_kwargs: Dict[str, float]
-) -> Callable[[hk.Params, jnp.ndarray, jnp.ndarray], jnp.ndarray]:
+) -> Callable[[hk.Params, jnp.ndarray, jnp.ndarray, jnp.ndarray, bool], jnp.ndarray]:
     try:
         loss_function = SUPPORTED_LOSSES[loss]
 
         @jax.jit
         def loss_function_wrapper(
-            params: hk.Params, x: jnp.ndarray, y_true: jnp.ndarray, rng: jnp.ndarray = None
+            params: hk.Params, x: jnp.ndarray, y_true: jnp.ndarray, rng: jnp.ndarray = None, is_training: bool = None
         ) -> jnp.ndarray:
             # rng argument can be used if net_transform.apply() is non-deterministic, and you require a "random seed"
-            y_pred: jnp.ndarray = net_transform.apply(params, rng, x)
+            try:
+                y_pred: jnp.ndarray = net_transform.apply(params, rng, x, is_training=is_training)
+            except TypeError:
+                y_pred: jnp.ndarray = net_transform.apply(params, rng, x)
             loss_value: jnp.ndarray = loss_function(y_true, y_pred, **loss_kwargs)  # type: ignore
             return loss_value