(params, state)
| 51 | |
| 52 | # Initialize the optimizer state to match the parameter state |
| 53 | def update_state(params, state): |
| 54 | if isinstance(params, (list, tuple)): |
| 55 | state = list(state) |
| 56 | for i in range(len(state)): |
| 57 | state[i] = update_state(params[i], state[i]) |
| 58 | if len(state) != len(params): |
| 59 | state.extend(tree_map(lambda _: {}, params[len(state) :])) |
| 60 | return type(params)(state) |
| 61 | elif isinstance(params, dict): |
| 62 | for k, v in params.items(): |
| 63 | if k not in state: |
| 64 | state[k] = tree_map(lambda _: {}, v) |
| 65 | else: |
| 66 | state[k] = update_state(v, state[k]) |
| 67 | return state |
| 68 | else: |
| 69 | return state |
| 70 | |
| 71 | update_state(parameters, self._state) |
| 72 | tree_map(lambda p, s: s or self.init_single(p, s), parameters, self._state) |
nothing calls this directly
no test coverage detected