Change loss_function_wrapper(.) to use new dm-haiku v0.0.2 .apply structure

Author: asmith26

jax_toolkit/losses/tests/test_utils.py

@@ -21,7 +21,6 @@ def net_function(x: jnp.ndarray) -> jnp.ndarray:
             return predictions
 
         net_transform = hk.transform(net_function)
-        net_transform = hk.without_apply_rng(net_transform)  # as deterministic outcome
         actual_loss_function_wrapper = get_haiku_loss_function(net_transform, loss="mean_squared_error")
 
         # Check works

jax_toolkit/losses/utils.py

@@ -48,8 +48,11 @@ def get_haiku_loss_function(
         loss_function = SUPPORTED_LOSSES[loss]
 
         @jax.jit
-        def loss_function_wrapper(params: hk.Params, x: jnp.ndarray, y_true: jnp.ndarray) -> jnp.ndarray:
-            y_pred: jnp.ndarray = net_transform.apply(params, x)
+        def loss_function_wrapper(
+            params: hk.Params, x: jnp.ndarray, y_true: jnp.ndarray, rng: jnp.ndarray = None
+        ) -> jnp.ndarray:
+            # rng argument can be used is net_transform.apply() is non-deterministic, and you require and "random seed"
+            y_pred: jnp.ndarray = net_transform.apply(params, rng, x)
             loss_value: jnp.ndarray = loss_function(y_true, y_pred)
             return loss_value