Lays out an array tensor by tensor to prevent OOMs.
(self, arrays: Any, layouts: Any, xla_flags: dict[str, Any] | None = None)
| 176 | return x |
| 177 | |
| 178 | def _iterated_layout(self, arrays: Any, layouts: Any, xla_flags: dict[str, Any] | None = None) -> Any: |
| 179 | """Lays out an array tensor by tensor to prevent OOMs.""" |
| 180 | |
| 181 | def _layout(x, s, l): |
| 182 | if x.format == l: |
| 183 | return x |
| 184 | # Somehow this can be None sometimes. |
| 185 | dll = (l.layout if jax.__version_info__ >= (0, 6, 3) else l.device_local_layout) if isinstance(l, Format) else l |
| 186 | f = jax.jit(self._identity, out_shardings=Format(dll, s)).lower(x).compile(compiler_options=xla_flags) |
| 187 | y = f(x) |
| 188 | # Achieves donation of the input argument, but allows for different memory |
| 189 | # layouts and shapes. |
| 190 | jax.tree.map(lambda z: z.delete(), x) |
| 191 | jax.block_until_ready(y) |
| 192 | return y |
| 193 | |
| 194 | shardings = jax.tree.map(lambda x: x.sharding, arrays) |
| 195 | arrays = jax.tree.map(_layout, arrays, shardings, layouts) |
| 196 | return arrays |
| 197 | |
| 198 | def aot_compile( |
| 199 | self, params: Params, pass_rng_shape: bool, xla_flags: dict[str, Any] | None = None |