Wrapper to generate for ahead of time compilation.
(
self, params: Params, decode_state: DecodeState, rng: PRNGKeyType | None = None
)
| 141 | max_utils.print_cpu_ram_stats(label) |
| 142 | |
| 143 | def generate_aot( |
| 144 | self, params: Params, decode_state: DecodeState, rng: PRNGKeyType | None = None |
| 145 | ) -> tuple[DecodeState, engine_api.ResultTokens]: |
| 146 | """Wrapper to generate for ahead of time compilation.""" |
| 147 | |
| 148 | return self.generate(params=params, decode_state=decode_state, rng=rng) |
| 149 | |
| 150 | def _compile_generate_and_get_layouts( |
| 151 | self, params: Any, decode_state: Any, rng_shape: Any, xla_flags: dict[str, Any] | None = None |