| 471 | "state": ocp.PyTreeCheckpointHandler( |
| 472 | save_concurrent_gb=cfg.max_concurrent_save_gb, |
| 473 | restore_concurrent_gb=cfg.max_concurrent_restore_gb, |
| 474 | ), |
| 475 | }, |
| 476 | ) |
| 477 | |
| 478 | def _get_spec(self, *, step: int, state: Nested[Any]) -> Nested[Any]: |
| 479 | spec = {"index": [("step", step)]} |
| 480 | for path, value in utils.flatten_items(state): |
| 481 | if isinstance(value, (Tensor, TensorSpec)): |
| 482 | dtype = getattr(value.dtype, "dtype", value.dtype) |
| 483 | spec["index"].append( |
| 484 | (path, {"dtype": str(dtype), "shape": str(tuple(value.shape))}) |
| 485 | ) |
| 486 | elif isinstance(value, (tf.data.Iterator, PythonSavable)): |
| 487 | spec["index"].append((path, str(type(value)))) |
| 488 | else: |