jax.extend.xla.clear_hlo_module_transformation

jax.extend.xla.clear_hlo_module_transformation(name, stage=PipelineStage.PRE_SCHEDULER, platforms=None)

Clear a registered custom compiler pass.

Parameters:

• name (str) – The name of the compiler pass to clear.

• stage (PipelineStage) – The pipeline stage of the pass. Must be a PipelineStage enum.

• platforms (Sequence[str] | str | None) – The list of platforms to clear the pass for. If None, the pass is cleared for all known backends.

Returns:

True if the pass was found and cleared, False otherwise.

Return type:

bool