(self, learning_rate, weight_decay, multiply_by_parameter_scale)
| 266 | |
| 267 | @parameterized.parameters((0.1, 0, False), (0.1, 0.01, True), (0.1, 0.0, True)) |
| 268 | def test_adamw_optimizer(self, learning_rate, weight_decay, multiply_by_parameter_scale): |
| 269 | adam_update_transformation = None |
| 270 | if multiply_by_parameter_scale: |
| 271 | adam_update_transformation = scale_by_param_block_rms() |
| 272 | self._test_optimizer( |
| 273 | adamw_optimizer( |
| 274 | learning_rate=learning_rate, |
| 275 | b1=0.9, |
| 276 | b2=0.99, |
| 277 | eps=1e-5, |
| 278 | weight_decay=weight_decay, |
| 279 | adam_update_transformation=adam_update_transformation, |
| 280 | ) |
| 281 | ) |
| 282 | |
| 283 | @parameterized.parameters((0.1, 0, 0.5, False), (0.1, 0.01, 0.2, True), (0.1, 0.0, 0.3, True)) |
| 284 | def test_adamw_decoupled_optimizer( |
nothing calls this directly
no test coverage detected