def test_clip_by_global_norm(self): updates = self.per_step_updates for i in range(1, STEPS + 1): clipper = transform.clip_by_global_norm(1. / i) updates, _ = clipper.update(updates, None) # Check that the clipper actually works and global norm is <= max_norm self.assertAlmostEqual(transform.global_norm(updates), 1. / i, places=6) updates_step, _ = clipper.update(self.per_step_updates, None) # Check that continuously clipping won't cause numerical issues. chex.assert_tree_all_close(updates, updates_step, atol=1e-7, rtol=1e-7)
def test_scale_by_fromage(self): schedule = lambda c: 1.0 / (c + 1.0) fromage = transform.scale_by_fromage(step_size_factor_fn=schedule) params = self.init_params state = fromage.init(params) previous_norm = jnp.inf transform_fn = self.variant(fromage.update) for _ in range(STEPS): # Apply a step of fromage updates, state = transform_fn(self.per_step_updates, state, params) # Updates should get smaller due to the the learning schedule. norm = transform.global_norm(updates) self.assertLess(norm, previous_norm) previous_norm = norm