Copies OptStateSpec and optionally assigns with a different memory kind. Args: specs: Nested[OptStateSpec] to copy from. pattern: Regex to match the full path of each spec. Matched specs will have their memory kind replaced with `memory_kind`. memory_kind: Ne
(
specs: Nested[OptStateSpec],
*,
pattern: Union[None, str, re.Pattern] = None,
memory_kind: Optional[MemoryKind] = None,
)
| 143 | |
| 144 | |
| 145 | def copy_partition( |
| 146 | specs: Nested[OptStateSpec], |
| 147 | *, |
| 148 | pattern: Union[None, str, re.Pattern] = None, |
| 149 | memory_kind: Optional[MemoryKind] = None, |
| 150 | ) -> Nested[OptStateSpec]: |
| 151 | """Copies OptStateSpec and optionally assigns with a different memory kind. |
| 152 | |
| 153 | Args: |
| 154 | specs: Nested[OptStateSpec] to copy from. |
| 155 | pattern: Regex to match the full path of each spec. Matched specs will have their memory |
| 156 | kind replaced with `memory_kind`. |
| 157 | memory_kind: New memory kind. Default to None. |
| 158 | |
| 159 | Returns: |
| 160 | A Nested[OptStateSpec] with possibly a different memory kind. |
| 161 | """ |
| 162 | return jax.tree.map( |
| 163 | lambda path, spec: OptStateSpec( |
| 164 | dtype=spec.dtype, |
| 165 | shape=spec.shape, |
| 166 | mesh_axes=spec.mesh_axes, |
| 167 | memory_kind=( |
| 168 | memory_kind if pattern and re.fullmatch(pattern, path) else spec.memory_kind |
| 169 | ), |
| 170 | ), |
| 171 | tree_paths(specs), |
| 172 | specs, |
| 173 | ) |
| 174 | |
| 175 | |
| 176 | def trace_partition( |