Modularize the lowering context to give more control to the user. This can be the order of the output lowering, explicit input/outputs, optionalize the wrapping, control over aliasing, sharing the XLA builder object, and decouple the XLA build step.
Motivation
Enhance the control and give room for specific use cases - from functional (e.g. while loop signatures) to optimizations (e.g. re-using lowered tracings). The main objective is to give room for users to use the lowering context interchangeably with other components, namely the XLA builder.
Pitch
We can either extend the lowering context, add a wrapper over it that provides users with more control on the computation they want to generate, or implement a seamless integration with the XLA builder.
🚀 Feature
Modularize the lowering context to give more control to the user. This can be the order of the output lowering, explicit input/outputs, optionalize the wrapping, control over aliasing, sharing the XLA builder object, and decouple the XLA build step.
Motivation
Enhance the control and give room for specific use cases - from functional (e.g. while loop signatures) to optimizations (e.g. re-using lowered tracings). The main objective is to give room for users to use the lowering context interchangeably with other components, namely the XLA builder.
Pitch
We can either extend the lowering context, add a wrapper over it that provides users with more control on the computation they want to generate, or implement a seamless integration with the XLA builder.