using jax.numpy.where for masking

Author: activatedgeek

optax_swag/transform.py

@@ -20,8 +20,9 @@ def update_fn(updates: optax.Updates, state: SWAState,
     update_mask = next_step == 0
     n = state.n + 1 * update_mask
 
-    next_params = jax.tree_util.tree_map(lambda u, p: u * update_mask + p, updates, params)
-    next_mean = jax.tree_util.tree_map(lambda mu, np: (n * mu + np) / (n + 1), 
+    next_params = jax.tree_util.tree_map(lambda p, u: jnp.where(update_mask, p + u, p),
+                                         params, updates)
+    next_mean = jax.tree_util.tree_map(lambda mu, np: jnp.where(update_mask, (n * mu + np) / (n + 1), mu),
                                        state.mean, next_params)
 
     return updates, SWAState(step=next_step, n=n, mean=next_mean)
@@ -44,10 +45,11 @@ def update_fn(updates: optax.Updates, state: SWAGDiagState,
     update_mask = next_step == 0
     n = state.n + 1 * update_mask
 
-    next_params = jax.tree_util.tree_map(lambda u, p: u * update_mask + p, updates, params)
-    next_mean = jax.tree_util.tree_map(lambda mu, np: (n * mu + np) / (n + 1), 
+    next_params = jax.tree_util.tree_map(lambda p, u: jnp.where(update_mask, p + u, p),
+                                         params, updates)
+    next_mean = jax.tree_util.tree_map(lambda mu, np: jnp.where(update_mask, (n * mu + np) / (n + 1), mu),
                                        state.mean, next_params)
-    next_params2 = jax.tree_util.tree_map(lambda v, np: (n * v + jnp.square(np)) / (n + 1),
+    next_params2 = jax.tree_util.tree_map(lambda p2, np: jnp.where(update_mask, (n * p2 + jnp.square(np)) / (n + 1), p2),
                                           state.params2, next_params)
 
     return updates, SWAGDiagState(step=next_step, n=n, mean=next_mean, params2=next_params2)
@@ -74,13 +76,14 @@ def update_fn(updates: optax.Updates, state: SWAGState,
     update_mask = next_step == 0
     n = state.n + 1 * update_mask
 
-    next_params = jax.tree_util.tree_map(lambda u, p: u * update_mask + p, updates, params)
-    next_mean = jax.tree_util.tree_map(lambda mu, np: (n * mu + np) / (n + 1), 
+    next_params = jax.tree_util.tree_map(lambda p, u: jnp.where(update_mask, p + u, p),
+                                         params, updates)
+    next_mean = jax.tree_util.tree_map(lambda mu, np: jnp.where(update_mask, (n * mu + np) / (n + 1), mu),
                                        state.mean, next_params)
-    next_params2 = jax.tree_util.tree_map(lambda v, np: (n * v + jnp.square(np)) / (n + 1),
+    next_params2 = jax.tree_util.tree_map(lambda p2, np: jnp.where(update_mask, (n * p2 + jnp.square(np)) / (n + 1), p2),
                                           state.params2, next_params)
-    next_dparams = jax.tree_util.tree_map(lambda np, nmu, dp: jnp.where(update_mask, dp.at[state.c].set(np - nmu), dp),
-                                          next_params, next_mean, state.dparams)
+    next_dparams = jax.tree_util.tree_map(lambda dp, np, nmu: jnp.where(update_mask, dp.at[state.c].set(np - nmu), dp),
+                                          state.dparams, next_params, next_mean)
 
     c = (state.c + 1 * update_mask) % rank