A dummy class implementing PythonSavable.
| 605 | if not test_utils.is_supported_mesh_shape(mesh_shape): |
| 606 | return |
| 607 | |
| 608 | class _DummySavable: |
| 609 | """A dummy class implementing PythonSavable.""" |
| 610 | |
| 611 | def __init__(self, max_step: int): |
| 612 | self._max_step = max_step |
| 613 | self._state = dict(step=0, values=b"123") |
| 614 | |
| 615 | def get_state(self): |
| 616 | return self._state |
| 617 | |
| 618 | def set_state(self, state): |
| 619 | self._state = state |
| 620 | |
| 621 | def __iter__(self): |
| 622 | for i in range(self._state["step"], self._max_step): |
| 623 | self._state["step"] += 1 |
| 624 | yield i |
| 625 | |
| 626 | with _mesh(mesh_shape): |
no outgoing calls