Fix uninitialized axis_env error when MLIR lowering is disabled

PiperOrigin-RevId: 429267926
abaelhe · Feb 17, 2022 · 57f4232 · 57f4232
1 parent 15295a8
commit 57f4232
Showing 1 changed file with 1 addition and 0 deletions.
diff --git a/jax/interpreters/pxla.py b/jax/interpreters/pxla.py
@@ -2173,6 +2173,7 @@ def lower_mesh_computation(
     out_partitions_t = xla.tuple_sharding_proto(out_partitions)
     partitions_proto = True
     axis_ctx = mlir.SPMDAxisContext(mesh)
+    axis_env = axis_ctx.axis_env
   else:
     replicated_args = [not axis for axis in in_axes]
     in_partitions = None