Esempio n. 1
0
 def test_trace_should_keep_track_of_momentum(self):
     params = np.zeros(6)
     updates = np.ones(6)
     opt = state.init(optix.trace(0.99, False))(random.PRNGKey(0), params,
                                                updates)
     onp.testing.assert_array_equal(opt.trace, np.zeros(6))
     opt = opt.update(params, updates)
     onp.testing.assert_array_equal(opt.trace, np.ones(6))
     onp.testing.assert_array_equal(opt(params, updates), 1.99 * np.ones(6))
Esempio n. 2
0
 def test_trace_should_keep_track_of_momentum_with_nesterov(self):
   params = jnp.zeros(6)
   updates = jnp.ones(6)
   opt = state.init(optix.trace(0.99, True))(random.PRNGKey(0),
                                             updates, params)
   np.testing.assert_array_equal(opt.trace, jnp.zeros(6))
   opt = opt.update(updates, params)
   np.testing.assert_array_equal(opt.trace, jnp.ones(6))
   np.testing.assert_array_equal(
       opt(updates, params), (1.99 + 0.99**2) * jnp.ones(6))