(self, checkpointer_cls, custom_dict_type)
| 423 | custom_dict_type=(utils.VDict,), |
| 424 | ) |
| 425 | def test_custom_dict(self, checkpointer_cls, custom_dict_type): |
| 426 | mesh_shape = (1, 1) |
| 427 | if not test_utils.is_supported_mesh_shape(mesh_shape): |
| 428 | return |
| 429 | with _mesh(mesh_shape): |
| 430 | cfg = _checkpointer_config(checkpointer_cls) |
| 431 | ckpt: Checkpointer = cfg.instantiate(parent=None) |
| 432 | state0 = custom_dict_type( |
| 433 | x=jnp.zeros([], dtype=jnp.int32), y=jnp.ones([2], dtype=jnp.float32) |
| 434 | ) |
| 435 | |
| 436 | ckpt.save(step=100, state=state0) |
| 437 | ckpt.wait_until_finished() |
| 438 | |
| 439 | # Restore with state structure hints will preserve VDict. |
| 440 | step, restored_state = ckpt.restore(step=None, state=state0) |
| 441 | self.assertEqual(100, step) |
| 442 | self.assertEqual(type(restored_state), custom_dict_type) |
| 443 | self.assertIn( |
| 444 | custom_dict_type.__name__, str(jax.tree_util.tree_structure(restored_state)) |
| 445 | ) |
| 446 | self.assertNestedEqual(state0, restored_state) |
| 447 | ckpt.stop() |
| 448 | |
| 449 | def test_elastic_input_iterator(self): |
| 450 | # In this test case, we want to simulate the training with |
nothing calls this directly
no test coverage detected