Seeking Guidance on Long-Term Support for Custom Backend Integration via stablehlo.CustomCallOp

[liangyunwang]

Hello JAX team and community,

We are integrating a custom compiler backend with JAX and have a question about the long-term stability of our proposed integration pattern. Our goal is to lower a JAX primitive to a stablehlo.CustomCallOp that points to a pre-compiled module for our backend.

Our current strategy involves a JAX primitive that lowers to a stablehlo.CustomCallOp. We use the backend_config attribute to pass a file path to our serialized custom module and output_operand_aliases to manage in-place buffer updates.

Our core concern is the long-term stability and official support for this specific pattern.

1. Is stablehlo.CustomCallOp the intended, long-term supported mechanism for integrating external compilers that consume MLIR from JAX?
2. Is using the backend_config attribute to pass a file path (or a large serialized program string) a robust and future-proof pattern? We've seen it used for simple attributes, but our use case involves passing an entire compiled artifact reference.
3. Is the output_operand_aliases attribute a stable part of the API for custom calls? This is critical for our performance as it allows us to perform in-place operations.
4. Are there any alternative integration paths or upcoming JAX features we should consider that might offer better long-term stability for this kind of backend integration?
5. As we build our infrastructure around this pattern, getting clarity on its long-term viability is crucial for us to ensure our backend remains compatible with future JAX releases.

Thank you for your time and any guidance you can provide!