Mock device for testing.
| 1599 | |
| 1600 | def test_two_per_param_scales(self): |
| 1601 | def config_fn(): |
| 1602 | trainer_cfg = SpmdTrainer.default_config() |
| 1603 | trainer_cfg.model = Linear.default_config().set( |
| 1604 | name="test_model", input_dim=3, output_dim=2 |
| 1605 | ) |
| 1606 | l2_per_param_scale = config_for_function(optimizers.per_param_scale_by_path).set( |
| 1607 | description="l2_regularizer_scale", |
| 1608 | scale_by_path=[ |
| 1609 | (".*bias.*", 0), |
no outgoing calls