Update jax.tree_map to jax.tree_util.tree_map.
@@ -113,10 +113,10 @@ class SGDLearner(acme.Learner):
return new_training_state, extra
def postprocess_aux(extra: LossExtra) -> LossExtra:
- reverb_update = jax.tree_map(lambda a: jnp.reshape(a, (-1, *a.shape[2:])),
- extra.reverb_update)
+ reverb_update = jax.tree_util.tree_map(
+ lambda a: jnp.reshape(a, (-1, *a.shape[2:])), extra.reverb_update)
return extra._replace(
- metrics=jax.tree_map(jnp.mean, extra.metrics),
+ metrics=jax.tree_util.tree_map(jnp.mean, extra.metrics),
reverb_update=reverb_update)
self._num_sgd_steps_per_step = num_sgd_steps_per_step
