Ahead of time compilation of generate with auto layout, relayout parameters.
(
self, params: Params, pass_rng_shape: bool, xla_flags: dict[str, Any] | None = None
)
| 196 | return arrays |
| 197 | |
| 198 | def aot_compile( |
| 199 | self, params: Params, pass_rng_shape: bool, xla_flags: dict[str, Any] | None = None |
| 200 | ) -> tuple[Any, Params, Any]: |
| 201 | """Ahead of time compilation of generate with auto layout, relayout parameters.""" |
| 202 | if pass_rng_shape: |
| 203 | rng_shape = jax.ShapeDtypeStruct([4], jax.numpy.dtype("uint32")) |
| 204 | else: |
| 205 | rng_shape = None |
| 206 | self.decode_state_shapes = jax.eval_shape(self.init_decode_state, rng_shape) |
| 207 | |
| 208 | generate_executable, self.param_layouts, _, self.decode_state_layouts = self._compile_generate_and_get_layouts( |
| 209 | self.abstract_params, self.decode_state_shapes, rng_shape, xla_flags |
| 210 | ) |
| 211 | return ( |
| 212 | generate_executable, |
| 213 | self._iterated_layout(params, self.param_layouts), |
| 214 | jax.jit(self.init_decode_state, in_shardings=(None), out_shardings=self.decode_state_layouts) |
| 215 | .lower(rng_shape) |
| 216 | .compile(), |
| 217 | ) |
| 218 | |
| 219 | def load_params(self, *args, params=None, rng: PRNGKeyType | None = None, **kwargs) -> Params: |
| 220 | """Load Parameters from GCS or reshard given Parameters""" |
no test coverage detected