(self)
| 46 | """Tests initialization of training and decode states in max_utils.py""" |
| 47 | |
| 48 | def setUp(self): |
| 49 | self.model = nn.Dense(features=5) |
| 50 | self.key1, self.key2 = random.split(random.key(0)) |
| 51 | self.input = random.normal(self.key1, (10,)) # Dummy input data |
| 52 | self.params = self.model.init(self.key2, self.input) |
| 53 | |
| 54 | def test_calculate_num_params_from_pytree(self): |
| 55 | example_tree = [ |