jax.extend.xla module

XLA compiler extensions.

PipelineStage(value[, names, module, ...])
clear_hlo_module_transformation(name[, ...])	Clear a registered custom compiler pass.
register_hlo_module_transformation(callback, ...)	Register a custom compiler pass that transforms HLO modules.