mirror of
https://github.com/ROCm/jax.git
synced 2025-04-18 12:56:07 +00:00
5cbc38d4f5
Previously, we kept the `dim_vars` in the `mlir.ModuleContext`. Now we replace that with a mutable `ShapePolyLoweringState` that also tracks whether we encounter shape polymorphism anywhere in the lowering. For this purpose, we also add `shape_poly_state` to the lowering.compile_args. We need to keep track of whether a module contains dimension variables because such modules need shape refinement before they can be converted to MHLO and compiled. For now, we just test that we set the `Exported.module_uses_dim_vars` correctly.