Wraps a lowered checkified function.
| 21 | |
| 22 | |
| 23 | class _CheckifyLoweredFnWrapper: |
| 24 | """Wraps a lowered checkified function.""" |
| 25 | |
| 26 | def __init__( |
| 27 | self, |
| 28 | lowered: jax.stages.Lowered, |
| 29 | ): |
| 30 | self._lowered = lowered |
| 31 | |
| 32 | def compile( |
| 33 | self, compiler_options: Optional[dict[str, Union[str, bool]]] = None |
| 34 | ) -> _CheckifyCompiledFnWrapper: |
| 35 | """Compile the function with provided options.""" |
| 36 | compiled = self._lowered.compile(compiler_options=compiler_options) |
| 37 | return _CheckifyCompiledFnWrapper(compiled) |
| 38 | |
| 39 | |
| 40 | class CheckifyJitFnWrapper: |